Refactor switch support in NodePath#getCompletionRecords (#13030)

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/condition-before-break/output.js

@@ -1,6 +1,9 @@
 const x = n => function () {
   switch (n) {
     case 0:
-      if (true) return void 0;
+      if (true) {
+        return void 0;
+      }
+
   }
 }();

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/declaration-before-break/exec.js

@@ -0,0 +1,38 @@
+const x = (n) => do {
+  switch (n) {
+    case 0:
+    case 6:
+      const b = 1;
+      break;
+    case 1: {
+      ("a");
+      {
+        const c = 1;
+        {
+          break;
+        }
+      }
+    }
+    case 2:
+    case 3: {
+      ("b");
+      if (n === 2) {
+        const c = 1;
+      } else {
+        ("c");
+      }
+      {
+        break;
+      }
+    }
+    default:
+      "bar";
+  }
+};
+
+expect(x(0)).toBeUndefined();
+expect(x(1)).toBeUndefined();
+expect(x(2)).toBeUndefined();
+expect(x(3)).toBe("c");
+expect(x(6)).toBeUndefined();
+expect(x(7)).toBe("bar");

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/declaration-before-break/input.js

@@ -4,6 +4,28 @@ const x = (n) => do {
     case 6:
       const b = 1;
       break;
-    default: 'bar';
+    case 1: {
+      ("a");
+      {
+        const c = 1;
+        {
+          break;
         }
-}
+      }
+    }
+    case 2:
+    case 3: {
+      ("b");
+      if (n === 2) {
+        const c = 1;
+      } else {
+        ("c");
+      }
+      {
+        break;
+      }
+    }
+    default:
+      "bar";
+  }
+};

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/declaration-before-break/output.js

@@ -5,7 +5,34 @@ const x = n => function () {
       const b = 1;
       return void 0;
 
+    case 1:
+      {
+        "a";
+        {
+          const c = 1;
+          {
+            return void 0;
+          }
+        }
+      }
+
+    case 2:
+    case 3:
+      {
+        "b";
+
+        if (n === 2) {
+          const c = 1;
+        } else {
+          return "c";
+        }
+
+        {
+          return void 0;
+        }
+      }
+
     default:
-      return 'bar';
+      return "bar";
   }
 }();

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/fallthrough-block/output.js

@@ -16,7 +16,9 @@ const x = n => function () {
       }
 
     case 3:
+      {
         return void 0;
+      }
 
     case 4:
       {

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/fallthrough-multiblock/output.js

@@ -0,0 +1,47 @@
+const x = n => function () {
+  switch (n) {
+    case 0:
+      {
+        "a";
+      }
+      {
+        "b";
+      }
+      ;
+
+    case 1:
+      {
+        return "c";
+      }
+      {
+        "d";
+      }
+      ;
+
+    case 2:
+      "a";
+      "b";
+
+    case 3:
+      {}
+      {
+        return void 0;
+      }
+
+    case 4:
+      {
+        "a";
+      }
+      {
+        "b";
+      }
+
+    case 5:
+      return "c";
+
+    case 6:
+      {}
+
+    case 7:
+  }
+}();

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/multi-consequent-multiblock/output.js

@@ -4,6 +4,7 @@ const x = n => function () {
       {
         return "a";
       }
+      {}
       {
         "b";
       }
@@ -20,6 +21,7 @@ const x = n => function () {
 
     case 2:
       return "a";
+      {}
       "b";
 
     case 3:
@@ -52,6 +54,9 @@ const x = n => function () {
       {
         return "a";
       }
+      {
+        "b";
+      }
       {
         break;
         "c";

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/nested-break-block/exec.js

@@ -0,0 +1,34 @@
+const x = (n) => do {
+  switch (n) {
+    case 0:
+      { "a"; { break; } }
+      { "b"; };
+    case 1:
+      { "c"; }
+      { "d"; { { { { break; }}}}};
+    case 2:
+    case 3:
+      { "e"; { if (true) { break; } } }
+      { break; }
+    case 4:
+      { "g"; { { break; "h"; } "i" } }
+    case 5:
+    case 6:
+      if (n === 5) {
+        "j"
+      } else {
+        "k"
+      }
+      { break; "l" }
+    case 7:
+  }
+}
+
+expect(x(0)).toBe('a')
+expect(x(1)).toBe('d')
+expect(x(2)).toBeUndefined()
+expect(x(3)).toBeUndefined()
+expect(x(4)).toBe("g")
+expect(x(5)).toBe("j")
+expect(x(6)).toBe("k")
+expect(x(7)).toBeUndefined()

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/nested-break-block/input.js

@@ -0,0 +1,25 @@
+const x = (n) => do {
+  switch (n) {
+    case 0:
+      { "a"; { break; } }
+      { "b"; };
+    case 1:
+      { "c"; }
+      { "d"; { { { { break; }}}}};
+    case 2:
+    case 3:
+      { "e"; { if (true) { break; } } }
+      { break; }
+    case 4:
+      { "g"; { { break; "h"; } "i" } }
+    case 5:
+    case 6:
+      if (n === 5) {
+        "j"
+      } else {
+        "k"
+      }
+      { break; "l" }
+    case 7:
+  }
+}

packages/babel-plugin-proposal-do-expressions/test/fixtures/do-expression-switch-case/nested-break-block/output.js

@@ -0,0 +1,66 @@
+const x = n => function () {
+  switch (n) {
+    case 0:
+      {
+        return "a";
+        {}
+      }
+      {
+        "b";
+      }
+      ;
+
+    case 1:
+      {
+        "c";
+      }
+      {
+        return "d";
+        {
+          {
+            {
+              {}
+            }
+          }
+        }
+      }
+      ;
+
+    case 2:
+    case 3:
+      {
+        "e";
+        {
+          if (true) {
+            return void 0;
+          }
+        }
+      }
+      {}
+
+    case 4:
+      {
+        return "g";
+        {
+          {
+            "h";
+          }
+          "i";
+        }
+      }
+
+    case 5:
+    case 6:
+      if (n === 5) {
+        return "j";
+      } else {
+        return "k";
+      }
+
+      {
+        "l";
+      }
+
+    case 7:
+  }
+}();

packages/babel-traverse/src/path/family.ts

@@ -4,6 +4,34 @@ import type TraversalContext from "../context";
 import NodePath from "./index";
 import * as t from "@babel/types";
 
+const NORMAL_COMPLETION = 0;
+const BREAK_COMPLETION = 1;
+
+type Completion = {
+  path: NodePath;
+  type: 0 | 1;
+};
+
+type CompletionContext = {
+  // whether the current context allows `break` statement. When it allows, we have
+  // to search all the statements for potential `break`
+  canHaveBreak: boolean;
+  // whether the statement is an immediate descendant of a switch case clause
+  inCaseClause: boolean;
+  // whether the `break` statement record should be populated to upper level
+  // when a `break` statement is an immediate descendant of a block statement, e.g.
+  // `{ break }`, it can influence the control flow in the upper levels.
+  shouldPopulateBreak: boolean;
+};
+
+function NormalCompletion(path: NodePath) {
+  return { type: NORMAL_COMPLETION, path };
+}
+
+function BreakCompletion(path: NodePath) {
+  return { type: BREAK_COMPLETION, path };
+}
+
 export function getOpposite(this: NodePath): NodePath | null {
   if (this.key === "left") {
     return this.getSibling("right");
@@ -15,114 +43,224 @@ export function getOpposite(this: NodePath): NodePath | null {
 
 function addCompletionRecords(
   path: NodePath | null | undefined,
+  records: Completion[],
+  context: CompletionContext,
+): Completion[] {
+  if (path) return records.concat(_getCompletionRecords(path, context));
+  return records;
+}
+
+function completionRecordForSwitch(
+  cases: NodePath<t.SwitchCase>[],
+  records: Completion[],
+  context: CompletionContext,
+): Completion[] {
+  // https://tc39.es/ecma262/#sec-runtime-semantics-caseblockevaluation
+  let lastNormalCompletions = [];
+  for (let i = 0; i < cases.length; i++) {
+    const casePath = cases[i];
+    const caseCompletions = _getCompletionRecords(casePath, context);
+    const normalCompletions = [];
+    const breakCompletions = [];
+    for (const c of caseCompletions) {
+      if (c.type === NORMAL_COMPLETION) {
+        normalCompletions.push(c);
+      }
+      if (c.type === BREAK_COMPLETION) {
+        breakCompletions.push(c);
+      }
+    }
+    if (normalCompletions.length) {
+      lastNormalCompletions = normalCompletions;
+    }
+    records = records.concat(breakCompletions);
+  }
+  records = records.concat(lastNormalCompletions);
+  return records;
+}
+
+function normalCompletionToBreak(completions: Completion[]) {
+  completions.forEach(c => {
+    c.type = BREAK_COMPLETION;
+  });
+}
+
+/**
+ * Determine how we should handle the break statement for break completions
+ *
+ * @param {Completion[]} completions
+ * @param {boolean} reachable Whether the break statement is reachable after
+   we mark the normal completions _before_ the given break completions as the final
+   completions. For example,
+   `{ 0 }; break;` is transformed to `{ return 0 }; break;`, the `break` here is unreachable
+   and thus can be removed without consequences. We may in the future reserve them instead since
+   we do not consistently remove unreachable statements _after_ break
+   `{ var x = 0 }; break;` is transformed to `{ var x = 0 }; return void 0;`, the `break` is reachable
+   because we can not wrap variable declaration under a return statement
+ */
+function replaceBreakStatementInBreakCompletion(
+  completions: Completion[],
+  reachable: boolean,
+) {
+  completions.forEach(c => {
+    if (c.path.isBreakStatement({ label: null })) {
+      if (reachable) {
+        c.path.replaceWith(t.unaryExpression("void", t.numericLiteral(0)));
+      } else {
+        c.path.remove();
+      }
+    }
+  });
+}
+
+function getStatementListCompletion(
   paths: NodePath[],
-): NodePath[] {
+  context: CompletionContext,
-  if (path) return paths.concat(path.getCompletionRecords());
+): Completion[] {
-  return paths;
+  let completions = [];
-}
+  if (context.canHaveBreak) {
-
+    let lastNormalCompletions = [];
-function findBreak(statements): NodePath | null {
+    for (let i = 0; i < paths.length; i++) {
-  let breakStatement;
+      const path = paths[i];
-  if (!Array.isArray(statements)) {
+      const newContext = { ...context, inCaseClause: false };
-    statements = [statements];
-  }
-
-  for (const statement of statements) {
       if (
-      statement.isDoExpression() ||
+        path.isBlockStatement() &&
-      statement.isProgram() ||
+        (context.inCaseClause || // case test: { break }
-      statement.isBlockStatement() ||
+          context.shouldPopulateBreak) // case test: { { break } }
-      statement.isCatchClause() ||
-      statement.isLabeledStatement()
       ) {
-      breakStatement = findBreak(statement.get("body"));
+        newContext.shouldPopulateBreak = true;
-    } else if (statement.isIfStatement()) {
-      breakStatement =
-        findBreak(statement.get("consequent")) ??
-        findBreak(statement.get("alternate"));
-    } else if (statement.isTryStatement()) {
-      breakStatement =
-        findBreak(statement.get("block")) ??
-        findBreak(statement.get("handler"));
-    } else if (statement.isBreakStatement()) {
-      breakStatement = statement;
-    }
-
-    if (breakStatement) {
-      return breakStatement;
-    }
-  }
-  return null;
-}
-
-function completionRecordForSwitch(cases, paths) {
-  let isLastCaseWithConsequent = true;
-
-  for (let i = cases.length - 1; i >= 0; i--) {
-    const switchCase = cases[i];
-    const consequent = switchCase.get("consequent");
-
-    let breakStatement = findBreak(consequent);
-
-    if (breakStatement) {
-      while (
-        breakStatement.key === 0 &&
-        breakStatement.parentPath.isBlockStatement()
-      ) {
-        breakStatement = breakStatement.parentPath;
-      }
-
-      const prevSibling = breakStatement.getPrevSibling();
-      if (
-        breakStatement.key > 0 &&
-        (prevSibling.isExpressionStatement() || prevSibling.isBlockStatement())
-      ) {
-        paths = addCompletionRecords(prevSibling, paths);
-        breakStatement.remove();
       } else {
-        breakStatement.replaceWith(breakStatement.scope.buildUndefinedNode());
+        newContext.shouldPopulateBreak = false;
-        paths = addCompletionRecords(breakStatement, paths);
       }
-    } else if (isLastCaseWithConsequent) {
+      const statementCompletions = _getCompletionRecords(path, newContext);
-      const statementFinder = statement =>
+      if (
-        !statement.isBlockStatement() ||
+        statementCompletions.length > 0 &&
-        statement.get("body").some(statementFinder);
+        // we can stop search `paths` when we have seen a `path` that is
-      const hasConsequent = consequent.some(statementFinder);
+        // effectively a `break` statement. Examples are
-      if (hasConsequent) {
+        // - `break`
-        paths = addCompletionRecords(consequent[consequent.length - 1], paths);
+        // - `if (true) { 1; break } else { 2; break }`
-        isLastCaseWithConsequent = false;
+        // - `{ break }```
+        // In other words, the paths after this `path` are unreachable
+        statementCompletions.every(c => c.type === BREAK_COMPLETION)
+      ) {
+        if (
+          lastNormalCompletions.length > 0 &&
+          statementCompletions.every(c =>
+            c.path.isBreakStatement({ label: null }),
+          )
+        ) {
+          // when a break completion has a path as BreakStatement, it must be `{ break }`
+          // whose completion value we can not determine, otherwise it would have been
+          // replaced by `replaceBreakStatementInBreakCompletion`
+          // When we have seen normal completions from the last statement
+          // it is safe to stop populating break and mark normal completions as break
+          normalCompletionToBreak(lastNormalCompletions);
+          completions = completions.concat(lastNormalCompletions);
+          // Declarations have empty completion record, however they can not be nested
+          // directly in return statement, i.e. `return (var a = 1)` is invalid.
+          if (lastNormalCompletions.some(c => c.path.isDeclaration())) {
+            completions = completions.concat(statementCompletions);
+            replaceBreakStatementInBreakCompletion(
+              statementCompletions,
+              /* reachable */ true,
+            );
+          }
+          replaceBreakStatementInBreakCompletion(
+            statementCompletions,
+            /* reachable */ false,
+          );
+        } else {
+          completions = completions.concat(statementCompletions);
+          if (!context.shouldPopulateBreak) {
+            replaceBreakStatementInBreakCompletion(
+              statementCompletions,
+              /* reachable */ true,
+            );
           }
         }
+        break;
       }
-  return paths;
+      if (i === paths.length - 1) {
+        completions = completions.concat(statementCompletions);
+      } else {
+        completions = completions.concat(
+          statementCompletions.filter(c => c.type === BREAK_COMPLETION),
+        );
+        lastNormalCompletions = statementCompletions.filter(
+          c => c.type === NORMAL_COMPLETION,
+        );
+      }
+    }
+  } else if (paths.length) {
+    // When we are in a context where `break` must not exist, we can skip linear
+    // search on statement lists and assume that the last statement determines
+    // the completion
+    completions = completions.concat(
+      _getCompletionRecords(paths[paths.length - 1], context),
+    );
+  }
+  return completions;
 }
 
+function _getCompletionRecords(
+  path: NodePath,
+  context: CompletionContext,
+): Completion[] {
+  let records = [];
+  if (path.isIfStatement()) {
+    records = addCompletionRecords(path.get("consequent"), records, context);
+    records = addCompletionRecords(path.get("alternate"), records, context);
+  } else if (path.isDoExpression() || path.isFor() || path.isWhile()) {
+    // @ts-expect-error(flow->ts): todo
+    records = addCompletionRecords(path.get("body"), records, context);
+  } else if (path.isProgram() || path.isBlockStatement()) {
+    records = records.concat(
+      // @ts-expect-error(flow->ts): todo
+      getStatementListCompletion(path.get("body"), context),
+    );
+  } else if (path.isFunction()) {
+    return _getCompletionRecords(path.get("body"), context);
+  } else if (path.isTryStatement()) {
+    records = addCompletionRecords(path.get("block"), records, context);
+    records = addCompletionRecords(path.get("handler"), records, context);
+  } else if (path.isCatchClause()) {
+    records = addCompletionRecords(path.get("body"), records, context);
+  } else if (path.isSwitchStatement()) {
+    records = completionRecordForSwitch(path.get("cases"), records, context);
+  } else if (path.isSwitchCase()) {
+    records = records.concat(
+      getStatementListCompletion(path.get("consequent"), {
+        canHaveBreak: true,
+        shouldPopulateBreak: false,
+        inCaseClause: true,
+      }),
+    );
+  } else if (path.isBreakStatement()) {
+    records.push(BreakCompletion(path));
+  } else {
+    records.push(NormalCompletion(path));
+  }
+
+  return records;
+}
+
+/**
+ * Retrieve the completion records of a given path.
+ * Note: to ensure proper support on `break` statement, this method
+ * will manipulate the AST around the break statement. Do not call the method
+ * twice for the same path.
+ *
+ * @export
+ * @param {NodePath} this
+ * @returns {NodePath[]} Completion records
+ */
 export function getCompletionRecords(this: NodePath): NodePath[] {
-  let paths = [];
+  const records = _getCompletionRecords(this, {
-
+    canHaveBreak: false,
-  if (this.isIfStatement()) {
+    shouldPopulateBreak: false,
-    paths = addCompletionRecords(this.get("consequent"), paths);
+    inCaseClause: false,
-    paths = addCompletionRecords(this.get("alternate"), paths);
+  });
-  } else if (this.isDoExpression() || this.isFor() || this.isWhile()) {
+  return records.map(r => r.path);
-    // @ts-expect-error(flow->ts): todo
-    paths = addCompletionRecords(this.get("body"), paths);
-  } else if (this.isProgram() || this.isBlockStatement()) {
-    // @ts-expect-error(flow->ts): todo
-    paths = addCompletionRecords(this.get("body").pop(), paths);
-  } else if (this.isFunction()) {
-    return this.get("body").getCompletionRecords();
-  } else if (this.isTryStatement()) {
-    paths = addCompletionRecords(this.get("block"), paths);
-    paths = addCompletionRecords(this.get("handler"), paths);
-  } else if (this.isCatchClause()) {
-    paths = addCompletionRecords(this.get("body"), paths);
-  } else if (this.isSwitchStatement()) {
-    paths = completionRecordForSwitch(this.get("cases"), paths);
-  } else {
-    paths.push(this);
-  }
-
-  return paths;
 }
 
 export function getSibling(this: NodePath, key: string | number): NodePath {