From 810264896dd85216a84a7fa29a691b6ceec0b1a9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Hu=C3=A1ng=20J=C3=B9nli=C3=A0ng?= <jlhwung@gmail.com>
Date: Fri, 3 Jan 2020 22:10:27 -0500
Subject: [PATCH] refactor: track [Return] parameter

---
 packages/babel-parser/src/parser/expression.js    |  5 +++++
 packages/babel-parser/src/parser/statement.js     |  2 +-
 .../babel-parser/src/util/production-parameter.js | 15 ++++++++++-----
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/packages/babel-parser/src/parser/expression.js b/packages/babel-parser/src/parser/expression.js
index 8d26c6b57d11..d469de858b04 100644
--- a/packages/babel-parser/src/parser/expression.js
+++ b/packages/babel-parser/src/parser/expression.js
@@ -43,6 +43,7 @@ import {
 import { ExpressionErrors } from "./util";
 import {
   PARAM_AWAIT,
+  PARAM_RETURN,
   PARAM,
   functionFlags,
 } from "../util/production-parameter";
@@ -1986,7 +1987,11 @@ export default class ExpressionParser extends LValParser {
         allowExpression,
         !oldStrict && useStrict,
       );
+      // FunctionBody[Yield, Await]:
+      //   StatementList[?Yield, ?Await, +Return] opt
+      this.prodParam.enter(this.prodParam.currentFlags() | PARAM_RETURN);
       node.body = this.parseBlock(true, false);
+      this.prodParam.exit();
       this.state.labels = oldLabels;
     }
 
diff --git a/packages/babel-parser/src/parser/statement.js b/packages/babel-parser/src/parser/statement.js
index 67f4279fe805..0a0d595a5457 100644
--- a/packages/babel-parser/src/parser/statement.js
+++ b/packages/babel-parser/src/parser/statement.js
@@ -575,7 +575,7 @@ export default class StatementParser extends ExpressionParser {
   }
 
   parseReturnStatement(node: N.ReturnStatement): N.ReturnStatement {
-    if (!this.scope.inFunction && !this.options.allowReturnOutsideFunction) {
+    if (!this.prodParam.hasReturn && !this.options.allowReturnOutsideFunction) {
       this.raise(this.state.start, "'return' outside of function");
     }
 
diff --git a/packages/babel-parser/src/util/production-parameter.js b/packages/babel-parser/src/util/production-parameter.js
index 9f468addc3d6..43644f6131ce 100644
--- a/packages/babel-parser/src/util/production-parameter.js
+++ b/packages/babel-parser/src/util/production-parameter.js
@@ -1,12 +1,13 @@
 // @flow
-export const PARAM = 0b00, // Initial Parameter flags
-  PARAM_YIELD = 0b01, // track [Await] production parameter
-  PARAM_AWAIT = 0b10; // track [Yield] production parameter
+export const PARAM = 0b000, // Initial Parameter flags
+  PARAM_YIELD = 0b001, // track [Yield] production parameter
+  PARAM_AWAIT = 0b010, // track [Await] production parameter
+  PARAM_RETURN = 0b100; // track [Return] production parameter
 
 // ProductionParameterHandler is a stack fashioned production parameter tracker
 // https://tc39.es/ecma262/#sec-grammar-notation
-// It only tracks [Await] and [Yield] parameter. The [In] parameter is tracked
-// in `noIn` argument of `parseExpression`.
+// The tracked parameters are defined above. Note that the [In] parameter is
+// tracked in `noIn` argument of `parseExpression`.
 //
 // Whenever [+Await]/[+Yield] appears in the right-hand sides of a production,
 // we must enter a new tracking stack. For example when parsing
@@ -48,6 +49,10 @@ export default class ProductionParameterHandler {
   get hasYield(): boolean {
     return (this.currentFlags() & PARAM_YIELD) > 0;
   }
+
+  get hasReturn(): boolean {
+    return (this.currentFlags() & PARAM_RETURN) > 0;
+  }
 }
 
 export function functionFlags(