From fbb10cd6b375da2120e3f538c814bdb58595db87 Mon Sep 17 00:00:00 2001
From: Daniel Rosenwasser <drosen@microsoft.com>
Date: Fri, 29 Aug 2014 14:18:13 -0700
Subject: [PATCH] Added getOccs support for return keywords.

---
 src/compiler/checker.ts  | 62 ++++++++++++++++++++--------------------
 src/compiler/types.ts    |  2 +-
 src/services/services.ts | 24 ++++++++++++++++
 3 files changed, 56 insertions(+), 32 deletions(-)

diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
index e1b2adde64f..575a9978f2f 100644
--- a/src/compiler/checker.ts
+++ b/src/compiler/checker.ts
@@ -4394,37 +4394,6 @@ module ts {
             return voidType;
         }
 
-        // WARNING: This has the same semantics as the forEach family of functions,
-        //          in that traversal terminates in the event that 'visitor' supplies a truthy value.
-        function forEachReturnStatement<T>(body: Block, visitor: (stmt: ReturnStatement) => T): T {
-
-            return traverse(body);
-
-            function traverse(node: Node): T {
-                switch (node.kind) {
-                    case SyntaxKind.ReturnStatement:
-                        return visitor(node);
-                    case SyntaxKind.Block:
-                    case SyntaxKind.FunctionBlock:
-                    case SyntaxKind.IfStatement:
-                    case SyntaxKind.DoStatement:
-                    case SyntaxKind.WhileStatement:
-                    case SyntaxKind.ForStatement:
-                    case SyntaxKind.ForInStatement:
-                    case SyntaxKind.WithStatement:
-                    case SyntaxKind.SwitchStatement:
-                    case SyntaxKind.CaseClause:
-                    case SyntaxKind.DefaultClause:
-                    case SyntaxKind.LabelledStatement:
-                    case SyntaxKind.TryStatement:
-                    case SyntaxKind.TryBlock:
-                    case SyntaxKind.CatchBlock:
-                    case SyntaxKind.FinallyBlock:
-                        return forEachChild(node, traverse);
-                }
-            }
-        }
-
         /// Returns a set of types relating to every return expression relating to a function block.
         function checkAndAggregateReturnExpressionTypes(body: Block, contextualMapper?: TypeMapper): Type[] {
             var aggregatedTypes: Type[] = [];
@@ -7155,4 +7124,35 @@ module ts {
 
         return checker;
     }
+
+    // WARNING: This has the same semantics as the forEach family of functions,
+    //          in that traversal terminates in the event that 'visitor' supplies a truthy value.
+    export function forEachReturnStatement<T>(body: Block, visitor: (stmt: ReturnStatement) => T): T {
+
+        return traverse(body);
+
+        function traverse(node: Node): T {
+            switch (node.kind) {
+                case SyntaxKind.ReturnStatement:
+                    return visitor(node);
+                case SyntaxKind.Block:
+                case SyntaxKind.FunctionBlock:
+                case SyntaxKind.IfStatement:
+                case SyntaxKind.DoStatement:
+                case SyntaxKind.WhileStatement:
+                case SyntaxKind.ForStatement:
+                case SyntaxKind.ForInStatement:
+                case SyntaxKind.WithStatement:
+                case SyntaxKind.SwitchStatement:
+                case SyntaxKind.CaseClause:
+                case SyntaxKind.DefaultClause:
+                case SyntaxKind.LabelledStatement:
+                case SyntaxKind.TryStatement:
+                case SyntaxKind.TryBlock:
+                case SyntaxKind.CatchBlock:
+                case SyntaxKind.FinallyBlock:
+                    return forEachChild(node, traverse);
+            }
+        }
+    }
 }
diff --git a/src/compiler/types.ts b/src/compiler/types.ts
index feccd3a298c..198182906f2 100644
--- a/src/compiler/types.ts
+++ b/src/compiler/types.ts
@@ -220,7 +220,7 @@ module ts {
         LastFutureReservedWord = YieldKeyword,
         FirstTypeNode = TypeReference,
         LastTypeNode = ArrayType,
-        FirstPunctuation= OpenBraceToken,
+        FirstPunctuation = OpenBraceToken,
         LastPunctuation = CaretEqualsToken
     }
 
diff --git a/src/services/services.ts b/src/services/services.ts
index e5bdd0bf751..0d75dbb2b36 100644
--- a/src/services/services.ts
+++ b/src/services/services.ts
@@ -2174,6 +2174,11 @@ module ts {
                         return getIfElseOccurrences(<IfStatement>node.parent);
                     }
                     break;
+                case SyntaxKind.ReturnKeyword:
+                    if (hasKind(node.parent, SyntaxKind.ReturnStatement)) {
+                        return getReturnOccurrences(<ReturnStatement>node.parent);
+                    }
+                    break;
                 case SyntaxKind.TryKeyword:
                 case SyntaxKind.CatchKeyword:
                 case SyntaxKind.FinallyKeyword:
@@ -2261,6 +2266,25 @@ module ts {
                 return result;
             }
 
+            function getReturnOccurrences(returnStatement: ReturnStatement): ReferenceEntry[]{
+                var node: Node = returnStatement;
+                while (!isAnyFunction(node) && node.parent) {
+                    node = node.parent;
+                }
+
+                // If we didn't find a containing function with a block body, bail out.
+                if (!(isAnyFunction(node) && hasKind((<FunctionDeclaration>node).body, SyntaxKind.FunctionBlock))) {
+                    return undefined;
+                }
+
+                var keywords: Node[] = []
+                forEachReturnStatement(<Block>(<FunctionDeclaration>node).body, returnStmt => {
+                    pushKeywordIf(keywords, returnStmt.getFirstToken(), SyntaxKind.ReturnKeyword);
+                });
+
+                return map(keywords, keywordToReferenceEntry);
+            }
+
             function getTryCatchFinallyOccurrences(tryStatement: TryStatement): ReferenceEntry[] {
                 var keywords: Node[] = [];