Merge pull request #593 from Microsoft/getOccurrencesReturn

Get occurrences for return keywords.

src/compiler/checker.ts

@@ -4401,37 +4401,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[] = [];
@@ -5812,17 +5781,6 @@ module ts {
             // TODO: Check that target label is valid
         }
 
-        function getContainingFunction(node: Node): SignatureDeclaration {
-            while (true) {
-                node = node.parent;
-                if (!node || node.kind === SyntaxKind.FunctionDeclaration || node.kind === SyntaxKind.FunctionExpression ||
-                    node.kind === SyntaxKind.ArrowFunction || node.kind === SyntaxKind.Method || node.kind === SyntaxKind.Constructor ||
-                    node.kind === SyntaxKind.GetAccessor || node.kind === SyntaxKind.SetAccessor) {
-                    return <SignatureDeclaration>node;
-                }
-            }
-        }
-
         function checkReturnStatement(node: ReturnStatement) {
             if (node.expression && !(getNodeLinks(node.expression).flags & NodeCheckFlags.TypeChecked)) {
                 var func = getContainingFunction(node);

src/compiler/parser.ts

@@ -350,6 +350,63 @@ module ts {
         }
     }
 
+    // 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);
+            }
+        }
+    }
+
+    export function isAnyFunction(node: Node): boolean {
+        if (node) {
+            switch (node.kind) {
+                case SyntaxKind.FunctionExpression:
+                case SyntaxKind.FunctionDeclaration:
+                case SyntaxKind.ArrowFunction:
+                case SyntaxKind.Method:
+                case SyntaxKind.GetAccessor:
+                case SyntaxKind.SetAccessor:
+                case SyntaxKind.Constructor:
+                    return true;
+            }
+        }
+
+        return false;
+    }
+
+    export function getContainingFunction(node: Node): SignatureDeclaration {
+        while (true) {
+            node = node.parent;
+            if (!node || isAnyFunction(node)) {
+                return <SignatureDeclaration>node;
+            }
+        }
+    }
+
     export function hasRestParameters(s: SignatureDeclaration): boolean {
         return s.parameters.length > 0 && (s.parameters[s.parameters.length - 1].flags & NodeFlags.Rest) !== 0;
     }

src/compiler/types.ts

@@ -220,7 +220,7 @@ module ts {
         LastFutureReservedWord = YieldKeyword,
         FirstTypeNode = TypeReference,
         LastTypeNode = ArrayType,
-        FirstPunctuation= OpenBraceToken,
+        FirstPunctuation = OpenBraceToken,
         LastPunctuation = CaretEqualsToken
     }
 

src/services/services.ts

@@ -1312,20 +1312,6 @@ module ts {
         return node.parent.kind === SyntaxKind.NewExpression && (<CallExpression>node.parent).func === node;
     }
 
-    function isAnyFunction(node: Node): boolean {
-        switch (node.kind) {
-            case SyntaxKind.FunctionExpression:
-            case SyntaxKind.FunctionDeclaration:
-            case SyntaxKind.ArrowFunction:
-            case SyntaxKind.Method:
-            case SyntaxKind.GetAccessor:
-            case SyntaxKind.SetAccessor:
-            case SyntaxKind.Constructor:
-                return true;
-        }
-        return false;
-    }
-
     function isNameOfFunctionDeclaration(node: Node): boolean {
         return node.kind === SyntaxKind.Identifier &&
             isAnyFunction(node.parent) && (<FunctionDeclaration>node.parent).name === node;
@@ -2171,6 +2157,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:
@@ -2258,6 +2249,22 @@ module ts {
                 return result;
             }
 
+            function getReturnOccurrences(returnStatement: ReturnStatement): ReferenceEntry[]{
+                var func = <FunctionDeclaration>getContainingFunction(returnStatement);
+
+                // If we didn't find a containing function with a block body, bail out.
+                if (!(func && hasKind(func.body, SyntaxKind.FunctionBlock))) {
+                    return undefined;
+                }
+
+                var keywords: Node[] = []
+                forEachReturnStatement(<Block>(<FunctionDeclaration>func).body, returnStatement => {
+                    pushKeywordIf(keywords, returnStatement.getFirstToken(), SyntaxKind.ReturnKeyword);
+                });
+
+                return map(keywords, keywordToReferenceEntry);
+            }
+
             function getTryCatchFinallyOccurrences(tryStatement: TryStatement): ReferenceEntry[] {
                 var keywords: Node[] = [];
 

tests/cases/fourslash/getOccurrencesReturn.ts

@@ -0,0 +1,33 @@
+/// <reference path='fourslash.ts' />
+
+////function f(a: number) {
+////    if (a > 0) {
+////        [|ret/**/urn|] (function () {
+////            return;
+////            return;
+////            return;
+////
+////            if (false) {
+////                return true;
+////            }
+////        })() || true;
+////    }
+////
+////    var unusued = [1, 2, 3, 4].map(x => { return 4 })
+////
+////    [|return|];
+////    [|return|] true;
+////}
+
+test.ranges().forEach(r => {
+    goTo.position(r.start);
+
+    test.ranges().forEach(range => {
+        verify.occurrencesAtPositionContains(range, false);
+    });
+});
+
+goTo.marker();
+test.ranges().forEach(range => {
+    verify.occurrencesAtPositionContains(range, false);
+});

tests/cases/fourslash/getOccurrencesReturn2.ts

@@ -0,0 +1,33 @@
+/// <reference path='fourslash.ts' />
+
+////function f(a: number) {
+////    if (a > 0) {
+////        return (function () {
+////            [|return|];
+////            [|ret/**/urn|];
+////            [|return|];
+////
+////            while (false) {
+////                [|return|] true;
+////            }
+////        })() || true;
+////    }
+////
+////    var unusued = [1, 2, 3, 4].map(x => { return 4 })
+////
+////    return;
+////    return true;
+////}
+
+test.ranges().forEach(r => {
+    goTo.position(r.start);
+
+    test.ranges().forEach(range => {
+        verify.occurrencesAtPositionContains(range, false);
+    });
+});
+
+goTo.marker();
+test.ranges().forEach(range => {
+    verify.occurrencesAtPositionContains(range, false);
+});

tests/cases/fourslash/getOccurrencesReturn3.ts

@@ -0,0 +1,28 @@
+/// <reference path='fourslash.ts' />
+
+////function f(a: number) {
+////    if (a > 0) {
+////        return (function () {
+////            return;
+////            return;
+////            return;
+////
+////            if (false) {
+////                return true;
+////            }
+////        })() || true;
+////    }
+////
+////    var unusued = [1, 2, 3, 4].map(x => { [|return|] 4 })
+////
+////    return;
+////    return true;
+////}
+
+test.ranges().forEach(r => {
+    goTo.position(r.start);
+
+    test.ranges().forEach(range => {
+        verify.occurrencesAtPositionContains(range, false);
+    });
+});

tests/cases/fourslash/getOccurrencesReturnBroken.ts

@@ -0,0 +1,54 @@
+/// <reference path='fourslash.ts' />
+
+////ret/*1*/urn;
+////retu/*2*/rn;
+////function f(a: number) {
+////    if (a > 0) {
+////        return (function () {
+////            () => [|return|];
+////            [|return|];
+////            [|return|];
+////
+////            if (false) {
+////                [|return|] true;
+////            }
+////        })() || true;
+////    }
+////
+////    var unusued = [1, 2, 3, 4].map(x => { return 4 })
+////
+////    return;
+////    return true;
+////}
+////
+////class A {
+////    ret/*3*/urn;
+////    r/*4*/eturn 8675309;
+////}
+
+// Note: For this test, these 'return's get highlighted as a result of a parse recovery 
+//       where if an arrow function starts with a statement, we try to parse a body
+//       as if it was missing curly braces. If the behavior changes in the future,
+//       a change to this test is very much welcome.
+test.ranges().forEach(r => {
+    goTo.position(r.start);
+
+    test.ranges().forEach(range => {
+        verify.occurrencesAtPositionContains(range, false);
+    });
+});
+
+for (var i = 1; i <= test.markers().length; i++) {
+    goTo.marker("" + i);
+
+    switch (i) {
+        case 0:
+        case 1:
+        case 4:
+            verify.occurrencesAtPositionCount(0);
+            break;
+        case 3:
+            verify.occurrencesAtPositionCount(1); // 'return' is an instance member
+            break;
+    }
+});

tests/cases/fourslash/getOccurrencesReturnNegatives.ts

@@ -0,0 +1,25 @@
+/// <reference path='fourslash.ts' />
+
+////function f(a: number) {
+////    if (a > 0) {
+////        return (function () {
+////            return/*1*/;
+////            return/*2*/;
+////            return/*3*/;
+////
+////            if (false) {
+////                return/*4*/ true;
+////            }
+////        })() || true;
+////    }
+////
+////    var unusued = [1, 2, 3, 4].map(x => { return/*5*/ 4 })
+////
+////    return/*6*/;
+////    return/*7*/ true;
+////}
+
+test.markers().forEach(m => {
+    goTo.position(m.position, m.fileName)
+    verify.occurrencesAtPositionCount(0);
+});