fix(51223): Go-to-definition for yield and await keyword; jump to res…

…pective function definition (#51838)

src/services/goToDefinition.ts

@@ -131,6 +131,18 @@ export function getDefinitionAtPosition(program: Program, sourceFile: SourceFile
         return functionDeclaration ? [createDefinitionFromSignatureDeclaration(typeChecker, functionDeclaration)] : undefined;
     }
 
+    if (node.kind === SyntaxKind.AwaitKeyword) {
+        const functionDeclaration = findAncestor(node, n => isFunctionLikeDeclaration(n)) as FunctionLikeDeclaration | undefined;
+        const isAsyncFunction = functionDeclaration && some(functionDeclaration.modifiers, (node) => node.kind === SyntaxKind.AsyncKeyword);
+        return isAsyncFunction ? [createDefinitionFromSignatureDeclaration(typeChecker, functionDeclaration)] : undefined;
+    }
+
+    if (node.kind === SyntaxKind.YieldKeyword) {
+        const functionDeclaration = findAncestor(node, n => isFunctionLikeDeclaration(n)) as FunctionLikeDeclaration | undefined;
+        const isGeneratorFunction = functionDeclaration && functionDeclaration.asteriskToken;
+        return isGeneratorFunction ? [createDefinitionFromSignatureDeclaration(typeChecker, functionDeclaration)] : undefined;
+    }
+
     if (isStaticModifier(node) && isClassStaticBlockDeclaration(node.parent)) {
         const classDecl = node.parent.parent;
         const { symbol, failedAliasResolution } = getSymbol(classDecl, typeChecker, stopAtAlias);

tests/cases/fourslash/goToDefinitionAwait1.ts

@@ -0,0 +1,12 @@
+/// <reference path="fourslash.ts" />
+
+//// async function /*end1*/foo() {
+////     [|/*start1*/await|] Promise.resolve(0);
+//// }
+
+//// function notAsync() {
+////     [|/*start2*/await|] Promise.resolve(0);
+//// }
+
+verify.goToDefinition("start1", "end1");
+verify.goToDefinition("start2", []);

tests/cases/fourslash/goToDefinitionAwait2.ts

@@ -0,0 +1,5 @@
+/// <reference path="fourslash.ts" />
+
+//// [|/*start*/await|] Promise.resolve(0);
+
+verify.goToDefinition("start", []);

tests/cases/fourslash/goToDefinitionAwait3.ts

@@ -0,0 +1,14 @@
+/// <reference path="fourslash.ts" />
+
+//// class C {
+////     notAsync() {
+////       [|/*start1*/await|] Promise.resolve(0);
+////     }
+////
+////     async /*end2*/foo() {
+////       [|/*start2*/await|] Promise.resolve(0);
+////     }
+//// }
+
+verify.goToDefinition("start1", []);
+verify.goToDefinition("start2", "end2");

tests/cases/fourslash/goToDefinitionAwait4.ts

@@ -0,0 +1,9 @@
+/// <reference path="fourslash.ts" />
+
+//// async function outerAsyncFun() {
+////     let /*end*/af = async () => {
+////       [|/*start*/await|] Promise.resolve(0);
+////     }
+//// }
+
+verify.goToDefinition("start", "end");

tests/cases/fourslash/goToDefinitionYield1.ts

@@ -0,0 +1,12 @@
+/// <reference path="fourslash.ts" />
+
+//// function* /*end1*/gen() {
+////     [|/*start1*/yield|] 0;
+//// }
+////
+//// const /*end2*/genFunction = function*() {
+////     [|/*start2*/yield|] 0;
+//// }
+
+verify.goToDefinition("start1", "end1");
+verify.goToDefinition("start2", "end2");

tests/cases/fourslash/goToDefinitionYield2.ts

@@ -0,0 +1,10 @@
+/// <reference path="fourslash.ts" />
+
+//// function* outerGen() {
+////     function* /*end*/gen() {
+////         [|/*start*/yield|] 0;
+////     }
+////     return gen
+//// }
+
+verify.goToDefinition("start", "end");

tests/cases/fourslash/goToDefinitionYield3.ts

@@ -0,0 +1,14 @@
+/// <reference path="fourslash.ts" />
+
+//// class C {
+////     notAGenerator() {
+////       [|/*start1*/yield|] 0;
+////     }
+////
+////     foo*/*end2*/() {
+////       [|/*start2*/yield|] 0;
+////     }
+//// }
+
+verify.goToDefinition("start1", []);
+verify.goToDefinition("start2", "end2");

tests/cases/fourslash/goToDefinitionYield4.ts

@@ -0,0 +1,7 @@
+/// <reference path="fourslash.ts" />
+
+//// function* gen() {
+////     class C { [/*start*/yield 10]() {} }
+//// }
+
+verify.goToDefinition("start", []);