Merge pull request #24431 from Zzzen/master

Document highlights on yield keywords highlight other occurrences in …

src/services/documentHighlights.ts

@@ -78,6 +78,8 @@ namespace ts.DocumentHighlights {
                 return useParent(node.parent, isAwaitExpression, getAsyncAndAwaitOccurrences);
             case SyntaxKind.AsyncKeyword:
                 return highlightSpans(getAsyncAndAwaitOccurrences(node));
+            case SyntaxKind.YieldKeyword:
+                return highlightSpans(getYieldOccurrences(node));
             default:
                 return isModifierKind(node.kind) && (isDeclaration(node.parent) || isVariableStatement(node.parent))
                     ? highlightSpans(getModifierOccurrences(node.kind, node.parent))
@@ -170,7 +172,7 @@ namespace ts.DocumentHighlights {
                     if (statement.kind === SyntaxKind.ContinueStatement) {
                         return false;
                     }
-                    // falls through
+                // falls through
                 case SyntaxKind.ForStatement:
                 case SyntaxKind.ForInStatement:
                 case SyntaxKind.ForOfStatement:
@@ -237,7 +239,7 @@ namespace ts.DocumentHighlights {
         }
     }
 
-    function pushKeywordIf(keywordList: Push<Node>, token: Node, ...expected: SyntaxKind[]): boolean {
+    function pushKeywordIf(keywordList: Push<Node>, token: Node | undefined, ...expected: SyntaxKind[]): boolean {
         if (token && contains(expected, token.kind)) {
             keywordList.push(token);
             return true;
@@ -384,18 +386,42 @@ namespace ts.DocumentHighlights {
             });
         }
 
-        forEachChild(func, aggregate);
+        forEachChild(func, child => {
+            traverseWithoutCrossingFunction(child, node => {
+                if (isAwaitExpression(node)) {
+                    pushKeywordIf(keywords, node.getFirstToken(), SyntaxKind.AwaitKeyword);
+                }
+            });
+        });
+
 
         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 getYieldOccurrences(node: Node): Node[] | undefined {
+        const func = getContainingFunction(node) as FunctionDeclaration;
+        if (!func) {
+            return undefined;
+        }
+
+        const keywords: Node[] = [];
+
+        forEachChild(func, child => {
+            traverseWithoutCrossingFunction(child, node => {
+                if (isYieldExpression(node)) {
+                    pushKeywordIf(keywords, node.getFirstToken(), SyntaxKind.YieldKeyword);
+                }
+            });
+        });
+
+        return keywords;
+    }
+
+    // Do not cross function/class/interface/module/type boundaries.
+    function traverseWithoutCrossingFunction(node: Node, cb: (node: Node) => void) {
+        cb(node);
+        if (!isFunctionLike(node) && !isClassLike(node) && !isInterfaceDeclaration(node) && !isModuleDeclaration(node) && !isTypeAliasDeclaration(node) && !isTypeNode(node)) {
+            forEachChild(node, child => traverseWithoutCrossingFunction(child, cb));
         }
     }
 

tests/cases/fourslash/getOccurrencesYield.ts

@@ -0,0 +1,22 @@
+/// <reference path='fourslash.ts' />
+
+////function* f() {
+//// [|yield|] 100;
+//// [|y/**/ield|] [|yield|] 200;
+////  class Foo {
+////      *memberFunction() {
+////          return yield 1;
+////      }
+////  }
+////  return function* g() {
+////    yield 1;
+////  }
+////}
+
+
+verify.rangesAreOccurrences(false);
+
+goTo.marker();
+for (const range of test.ranges()) {
+    verify.occurrencesAtPositionContains(range, false);
+}