diff --git a/src/services/documentHighlights.ts b/src/services/documentHighlights.ts index 69c75548fe537..65653483d0755 100644 --- a/src/services/documentHighlights.ts +++ b/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 = 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) && !isClassLike(node) && !isInterfaceDeclaration(node) && !isModuleDeclaration(node) && !isTypeAliasDeclaration(node) && !isTypeNode(node)) { + forEachChild(node, aggregate); + } + } + } + function getIfElseOccurrences(ifStatement: IfStatement, sourceFile: SourceFile): HighlightSpan[] { const keywords = getIfElseKeywords(ifStatement, sourceFile); const result: HighlightSpan[] = []; diff --git a/tests/cases/fourslash/getOccurrencesAsyncAwait.ts b/tests/cases/fourslash/getOccurrencesAsyncAwait.ts new file mode 100644 index 0000000000000..f337e755096ab --- /dev/null +++ b/tests/cases/fourslash/getOccurrencesAsyncAwait.ts @@ -0,0 +1,27 @@ +/// + +////[|async|] function f() { +//// [|await|] 100; +//// [|a/**/wait|] [|await|] 200; +////class Foo { +//// async memberFunction() { +//// await 1; +//// } +////} +//// 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); +} diff --git a/tests/cases/fourslash/getOccurrencesAsyncAwait2.ts b/tests/cases/fourslash/getOccurrencesAsyncAwait2.ts new file mode 100644 index 0000000000000..530f26bc18d37 --- /dev/null +++ b/tests/cases/fourslash/getOccurrencesAsyncAwait2.ts @@ -0,0 +1,16 @@ +/// + +////[|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); +} diff --git a/tests/cases/fourslash/getOccurrencesAsyncAwait3.ts b/tests/cases/fourslash/getOccurrencesAsyncAwait3.ts new file mode 100644 index 0000000000000..972b6b8b07d70 --- /dev/null +++ b/tests/cases/fourslash/getOccurrencesAsyncAwait3.ts @@ -0,0 +1,11 @@ +/// + +// Not valid TS ('await' expression is only allowed within an async function.) + +////a/**/wait 100; +////async function f() { +//// await 300; +////} + +goTo.marker(); +verify.occurrencesAtPositionCount(0); \ No newline at end of file