From dbe87a67415472c1f826802234c885cee41511ef Mon Sep 17 00:00:00 2001
From: Zzzen <843968788@qq.com>
Date: Sat, 19 May 2018 20:17:42 +0800
Subject: [PATCH] Document highlights on yield keywords highlight other
 occurrences in the same body

---
 src/services/documentHighlights.ts           | 48 +++++++++++++++-----
 tests/cases/fourslash/getOccurrencesYield.ts | 22 +++++++++
 2 files changed, 59 insertions(+), 11 deletions(-)
 create mode 100644 tests/cases/fourslash/getOccurrencesYield.ts

diff --git a/src/services/documentHighlights.ts b/src/services/documentHighlights.ts
index e6a78e098e9..cd28f46a841 100644
--- a/src/services/documentHighlights.ts
+++ b/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));
         }
     }
 
diff --git a/tests/cases/fourslash/getOccurrencesYield.ts b/tests/cases/fourslash/getOccurrencesYield.ts
new file mode 100644
index 00000000000..34a0b771ad3
--- /dev/null
+++ b/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);
+}