Document highlights on async/await keywords should highlight other occurrences in the same body fix #22944

src/services/documentHighlights.ts

@@ -74,6 +74,10 @@ namespace ts.DocumentHighlights {
             case SyntaxKind.GetKeyword:
             case SyntaxKind.SetKeyword:
                 return getFromAllDeclarations(isAccessor, [SyntaxKind.GetKeyword, SyntaxKind.SetKeyword]);
+            case SyntaxKind.AwaitKeyword:
+                return useParent(node.parent, isAwaitExpression, getAsyncAndAwaitOccurrences);
+            case SyntaxKind.AsyncKeyword:
+                return highlightSpans(getAsyncAndAwaitOccurrences(node));
             default:
                 return isModifierKind(node.kind) && (isDeclaration(node.parent) || isVariableStatement(node.parent))
                     ? highlightSpans(getModifierOccurrences(node.kind, node.parent))
@@ -368,6 +372,35 @@ namespace ts.DocumentHighlights {
         return keywords;
     }
 
+    function getAsyncAndAwaitOccurrences(node: Node): Node[] | undefined {
+        const func = <FunctionLikeDeclaration>getContainingFunction(node);
+        if (!func) {
+            return undefined;
+        }
+
+        const keywords: Node[] = [];
+
+        if (func.modifiers) {
+            func.modifiers.forEach(modifier => {
+                pushKeywordIf(keywords, modifier, SyntaxKind.AsyncKeyword);
+            });
+        }
+
+        forEachChild(func, aggregate);
+
+        return keywords;
+
+        function aggregate(node: Node): void {
+            if (isAwaitExpression(node)) {
+                pushKeywordIf(keywords, node.getFirstToken(), SyntaxKind.AwaitKeyword);
+            }
+            // Do not cross function boundaries.
+            if (!isFunctionLike(node)) {
+                forEachChild(node, aggregate);
+            }
+        }
+    }
+
     function getIfElseOccurrences(ifStatement: IfStatement, sourceFile: SourceFile): HighlightSpan[] {
         const keywords = getIfElseKeywords(ifStatement, sourceFile);
         const result: HighlightSpan[] = [];

tests/cases/fourslash/getOccurrencesAsyncAwait.ts

@@ -0,0 +1,22 @@
+/// <reference path='fourslash.ts' />
+
+////[|async|] function f() {
+//// [|await|] 100;
+//// [|a/**/wait|] [|await|] 200;
+//// return [|await|] async function () {
+////   await 300;
+//// }
+////}
+////async function g() {
+////    await 300;
+////    async function f() {
+////        await 400;
+////    }
+////}
+
+verify.rangesAreOccurrences(false);
+
+goTo.marker();
+for (const range of test.ranges()) {
+    verify.occurrencesAtPositionContains(range, false);
+}

tests/cases/fourslash/getOccurrencesAsyncAwait2.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////[|a/**/sync|] function f() {
+//// [|await|] 100;
+//// [|await|] [|await|] 200;
+//// return [|await|] async function () {
+////   await 300;
+//// }
+////}
+
+verify.rangesAreOccurrences(false);
+
+goTo.marker();
+for (const range of test.ranges()) {
+    verify.occurrencesAtPositionContains(range, false);
+}

tests/cases/fourslash/getOccurrencesAsyncAwait3.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+// Not valid TS ('await' expression is only allowed within an async function.)
+
+////a/**/wait 100;
+////async function f() {
+////    await 300;
+////}
+
+verify.rangesAreOccurrences(false);
+
+goTo.marker();
+verify.occurrencesAtPositionCount(0);