Fix several edge cases with context expression state (#8972)

* Fix several edge cases with context expression state

* Fix review comments

* Remove unused field

packages/babel-parser/src/tokenizer/context.js

@@ -12,7 +12,7 @@ export class TokContext {
     token: string,
     isExpr?: boolean,
     preserveSpace?: boolean,
-    override?: Function, // Takes a Tokenizer as a this-parameter, and returns void.
+    override?: ?Function, // Takes a Tokenizer as a this-parameter, and returns void.
   ) {
     this.token = token;
     this.isExpr = !!isExpr;
@@ -31,11 +31,12 @@ export const types: {
 } = {
   braceStatement: new TokContext("{", false),
   braceExpression: new TokContext("{", true),
-  templateQuasi: new TokContext("${", true),
+  templateQuasi: new TokContext("${", false),
   parenStatement: new TokContext("(", false),
   parenExpression: new TokContext("(", true),
   template: new TokContext("`", true, true, p => p.readTmplToken()),
   functionExpression: new TokContext("function", true),
+  functionStatement: new TokContext("function", false),
 };
 
 // Token-specific context update code
@@ -46,33 +47,26 @@ tt.parenR.updateContext = tt.braceR.updateContext = function() {
     return;
   }
 
-  const out = this.state.context.pop();
-  if (
-    out === types.braceStatement &&
-    this.curContext() === types.functionExpression
-  ) {
-    this.state.context.pop();
-    this.state.exprAllowed = false;
-  } else if (out === types.templateQuasi) {
-    this.state.exprAllowed = true;
-  } else {
-    this.state.exprAllowed = !out.isExpr;
+  let out = this.state.context.pop();
+  if (out === types.braceStatement && this.curContext().token === "function") {
+    out = this.state.context.pop();
   }
+
+  this.state.exprAllowed = !out.isExpr;
 };
 
 tt.name.updateContext = function(prevType) {
-  if (this.state.value === "of" && this.curContext() === types.parenStatement) {
-    this.state.exprAllowed = !prevType.beforeExpr;
-    return;
-  }
-
-  this.state.exprAllowed = false;
-
-  if (prevType === tt._let || prevType === tt._const || prevType === tt._var) {
-    if (lineBreak.test(this.input.slice(this.state.end))) {
-      this.state.exprAllowed = true;
+  let allowed = false;
+  if (prevType !== tt.dot) {
+    if (
+      (this.state.value === "of" && !this.state.exprAllowed) ||
+      (this.state.value === "yield" && this.state.inGenerator)
+    ) {
+      allowed = true;
     }
   }
+  this.state.exprAllowed = allowed;
+
   if (this.state.isIterator) {
     this.state.isIterator = false;
   }
@@ -107,8 +101,22 @@ tt.incDec.updateContext = function() {
 };
 
 tt._function.updateContext = tt._class.updateContext = function(prevType) {
-  if (this.state.exprAllowed && !this.braceIsBlock(prevType)) {
+  if (
+    prevType.beforeExpr &&
+    prevType !== tt.semi &&
+    prevType !== tt._else &&
+    !(
+      prevType === tt._return &&
+      lineBreak.test(this.input.slice(this.state.lastTokEnd, this.state.start))
+    ) &&
+    !(
+      (prevType === tt.colon || prevType === tt.braceL) &&
+      this.curContext() === types.b_stat
+    )
+  ) {
     this.state.context.push(types.functionExpression);
+  } else {
+    this.state.context.push(types.functionStatement);
   }
 
   this.state.exprAllowed = false;

packages/babel-parser/src/tokenizer/index.js

@@ -1324,14 +1324,25 @@ export default class Tokenizer extends LocationParser {
   }
 
   braceIsBlock(prevType: TokenType): boolean {
-    if (prevType === tt.colon) {
-      const parent = this.curContext();
-      if (parent === ct.braceStatement || parent === ct.braceExpression) {
-        return !parent.isExpr;
-      }
+    const parent = this.curContext();
+    if (parent === ct.functionExpression || parent === ct.functionStatement) {
+      return true;
+    }
+    if (
+      prevType === tt.colon &&
+      (parent === ct.braceStatement || parent === ct.braceExpression)
+    ) {
+      return !parent.isExpr;
     }
 
-    if (prevType === tt._return) {
+    // The check for `tt.name && exprAllowed` detects whether we are
+    // after a `yield` or `of` construct. See the `updateContext` for
+    // `tt.name`.
+    if (
+      prevType === tt._return ||
+      prevType === tt._yield ||
+      (prevType === tt.name && this.state.exprAllowed)
+    ) {
       return lineBreak.test(
         this.input.slice(this.state.lastTokEnd, this.state.start),
       );
@@ -1341,13 +1352,22 @@ export default class Tokenizer extends LocationParser {
       prevType === tt._else ||
       prevType === tt.semi ||
       prevType === tt.eof ||
-      prevType === tt.parenR
+      prevType === tt.parenR ||
+      prevType === tt.arrow
     ) {
       return true;
     }
 
     if (prevType === tt.braceL) {
-      return this.curContext() === ct.braceStatement;
+      return parent === ct.braceStatement;
+    }
+
+    if (
+      prevType === tt._var ||
+      prevType === tt._let ||
+      prevType === tt._const
+    ) {
+      return false;
     }
 
     if (prevType === tt.relational) {

packages/babel-parser/src/tokenizer/types.js

@@ -174,7 +174,7 @@ export const keywords = {
   new: new KeywordTokenType("new", { beforeExpr, startsExpr }),
   this: new KeywordTokenType("this", { startsExpr }),
   super: new KeywordTokenType("super", { startsExpr }),
-  class: new KeywordTokenType("class"),
+  class: new KeywordTokenType("class", { startsExpr }),
   extends: new KeywordTokenType("extends", { beforeExpr }),
   export: new KeywordTokenType("export"),
   import: new KeywordTokenType("import", { startsExpr }),