From d584737ea2afab9a42fc7ec44e4170e836e8e90b Mon Sep 17 00:00:00 2001
From: Yui T <yuisu@microsoft.com>
Date: Mon, 15 Dec 2014 14:47:09 -0800
Subject: [PATCH] Address code review

---
 src/compiler/checker.ts | 25 +++++++------------------
 1 file changed, 7 insertions(+), 18 deletions(-)

diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
index 972290095ed..d4376becd6a 100644
--- a/src/compiler/checker.ts
+++ b/src/compiler/checker.ts
@@ -8139,14 +8139,8 @@ module ts {
             // Grammar checking
             var parent = node.parent;
             var inFunctionBlock = false;
-            while (parent && parent.kind !== SyntaxKind.SourceFile) {
-                inFunctionBlock = isFunctionBlock(parent);
-                if (inFunctionBlock) {
-                    break;
-                }
-                parent = parent.parent;
-            }
-            if (!inFunctionBlock) {
+            var functionBlock = getContainingFunction(node);
+            if (!functionBlock) {
                 grammarErrorOnFirstToken(node, Diagnostics.A_return_statement_can_only_be_used_within_a_function_body);
             }
 
@@ -8174,21 +8168,16 @@ module ts {
 
         function checkWithStatement(node: WithStatement) {
             // Grammar checking for withStatement
+            var hasStrictModeError = false;
             if (node.parserContextFlags & ParserContextFlags.StrictMode) {
                 grammarErrorOnFirstToken(node, Diagnostics.with_statements_are_not_allowed_in_strict_mode);
+                hasStrictModeError = true;
             }
 
             // Grammar checking for invalid use of return statement
-            if (node.statement.kind === SyntaxKind.ReturnStatement) {
-                // Grammar check for invalid use of return statement
-                grammarErrorOnFirstToken(node.statement, Diagnostics.A_return_statement_can_only_be_used_within_a_function_body);
-            }
-            else if (node.statement.kind === SyntaxKind.Block) {
-                forEach((<Block>node.statement).statements, statement => {
-                    if (statement.kind === SyntaxKind.ReturnStatement) {
-                        // Grammar check for invalid use of return statement
-                        grammarErrorOnFirstToken(node.statement, Diagnostics.A_return_statement_can_only_be_used_within_a_function_body);
-                    }
+            if (!hasStrictModeError) {
+                forEachReturnStatement(<Block>node.statement, (returnStatement) => {
+                    grammarErrorOnFirstToken(returnStatement, Diagnostics.A_return_statement_can_only_be_used_within_a_function_body);
                 });
             }