[ts5.0] Better inlining of constants in enums (#15379)

liuxingbaoyu · web-flow · commit 52ccb069988f · 2023-02-18T16:00:40.000+01:00
diff --git a/packages/babel-plugin-transform-typescript/src/const-enum.ts b/packages/babel-plugin-transform-typescript/src/const-enum.ts
@@ -27,7 +27,7 @@ export default function transpileConstEnum(
     );
   }
 
-  const entries = translateEnumValues(path, t);
+  const { enumValues: entries } = translateEnumValues(path, t);
 
   if (isExported) {
     const obj = t.objectExpression(
diff --git a/packages/babel-plugin-transform-typescript/src/enum.ts b/packages/babel-plugin-transform-typescript/src/enum.ts
@@ -1,10 +1,11 @@
-import { template } from "@babel/core";
+import { template, types as t } from "@babel/core";
 import type { NodePath } from "@babel/traverse";
-import type * as t from "@babel/types";
 import assert from "assert";
 
 type t = typeof t;
 
+const ENUMS = new WeakMap<t.Identifier, PreviousEnumMembers>();
+
 export default function transpileEnum(
   path: NodePath<t.TSEnumDeclaration>,
   t: t,
@@ -17,7 +18,7 @@ export default function transpileEnum(
   }
 
   const name = node.id.name;
-  const fill = enumFill(path, t, node.id);
+  const { wrapper: fill, data } = enumFill(path, t, node.id);
 
   switch (path.parent.type) {
     case "BlockStatement":
@@ -33,6 +34,7 @@ export default function transpileEnum(
         path.scope.registerDeclaration(
           path.replaceWith(makeVar(node.id, t, isGlobal ? "var" : "let"))[0],
         );
+        ENUMS.set(path.scope.getBindingIdentifier(name), data);
       }
       break;
     }
@@ -81,7 +83,7 @@ const buildEnumMember = (isString: boolean, options: Record<string, unknown>) =>
  * `(function (E) { ... assignments ... })(E || (E = {}));`
  */
 function enumFill(path: NodePath<t.TSEnumDeclaration>, t: t, id: t.Identifier) {
-  const x = translateEnumValues(path, t);
+  const { enumValues: x, data } = translateEnumValues(path, t);
   const assignments = x.map(([memberName, memberValue]) =>
     buildEnumMember(t.isStringLiteral(memberValue), {
       ENUM: t.cloneNode(id),
@@ -90,10 +92,13 @@ function enumFill(path: NodePath<t.TSEnumDeclaration>, t: t, id: t.Identifier) {
     }),
   );
 
-  return buildEnumWrapper({
-    ID: t.cloneNode(id),
-    ASSIGNMENTS: assignments,
-  });
+  return {
+    wrapper: buildEnumWrapper({
+      ID: t.cloneNode(id),
+      ASSIGNMENTS: assignments,
+    }),
+    data: data,
+  };
 }
 
 /**
@@ -131,109 +136,172 @@ const enumSelfReferenceVisitor = {
   ReferencedIdentifier,
 };
 
-export function translateEnumValues(
-  path: NodePath<t.TSEnumDeclaration>,
-  t: t,
-): Array<[name: string, value: t.Expression]> {
+export function translateEnumValues(path: NodePath<t.TSEnumDeclaration>, t: t) {
   const seen: PreviousEnumMembers = new Map();
   // Start at -1 so the first enum member is its increment, 0.
   let constValue: number | string | undefined = -1;
   let lastName: string;
 
-  return path.get("members").map(memberPath => {
-    const member = memberPath.node;
-    const name = t.isIdentifier(member.id) ? member.id.name : member.id.value;
-    const initializer = member.initializer;
-    let value: t.Expression;
-    if (initializer) {
-      constValue = computeConstantValue(initializer, seen);
-      if (constValue !== undefined) {
-        seen.set(name, constValue);
-        if (typeof constValue === "number") {
-          value = t.numericLiteral(constValue);
+  return {
+    data: seen,
+    enumValues: path.get("members").map(memberPath => {
+      const member = memberPath.node;
+      const name = t.isIdentifier(member.id) ? member.id.name : member.id.value;
+      const initializerPath = memberPath.get("initializer");
+      const initializer = member.initializer;
+      let value: t.Expression;
+      if (initializer) {
+        constValue = computeConstantValue(initializerPath, seen);
+        if (constValue !== undefined) {
+          seen.set(name, constValue);
+          if (typeof constValue === "number") {
+            value = t.numericLiteral(constValue);
+          } else {
+            assert(typeof constValue === "string");
+            value = t.stringLiteral(constValue);
+          }
         } else {
-          assert(typeof constValue === "string");
-          value = t.stringLiteral(constValue);
+          if (initializerPath.isReferencedIdentifier()) {
+            ReferencedIdentifier(initializerPath, {
+              t,
+              seen,
+              path,
+            });
+          } else {
+            initializerPath.traverse(enumSelfReferenceVisitor, {
+              t,
+              seen,
+              path,
+            });
+          }
+
+          value = initializerPath.node;
+          seen.set(name, undefined);
         }
+      } else if (typeof constValue === "number") {
+        constValue += 1;
+        value = t.numericLiteral(constValue);
+        seen.set(name, constValue);
+      } else if (typeof constValue === "string") {
+        throw path.buildCodeFrameError("Enum member must have initializer.");
       } else {
-        const initializerPath = memberPath.get("initializer");
-
-        if (initializerPath.isReferencedIdentifier()) {
-          ReferencedIdentifier(initializerPath, {
-            t,
-            seen,
-            path,
-          });
-        } else {
-          initializerPath.traverse(enumSelfReferenceVisitor, { t, seen, path });
-        }
-
-        value = initializerPath.node;
+        // create dynamic initializer: 1 + ENUM["PREVIOUS"]
+        const lastRef = t.memberExpression(
+          t.cloneNode(path.node.id),
+          t.stringLiteral(lastName),
+          true,
+        );
+        value = t.binaryExpression("+", t.numericLiteral(1), lastRef);
         seen.set(name, undefined);
       }
-    } else if (typeof constValue === "number") {
-      constValue += 1;
-      value = t.numericLiteral(constValue);
-      seen.set(name, constValue);
-    } else if (typeof constValue === "string") {
-      throw path.buildCodeFrameError("Enum member must have initializer.");
-    } else {
-      // create dynamic initializer: 1 + ENUM["PREVIOUS"]
-      const lastRef = t.memberExpression(
-        t.cloneNode(path.node.id),
-        t.stringLiteral(lastName),
-        true,
-      );
-      value = t.binaryExpression("+", t.numericLiteral(1), lastRef);
-      seen.set(name, undefined);
-    }
 
-    lastName = name;
-    return [name, value];
-  });
+      lastName = name;
+      return [name, value];
+    }) as Array<[name: string, value: t.Expression]>,
+  };
 }
 
 // Based on the TypeScript repository's `computeConstantValue` in `checker.ts`.
 function computeConstantValue(
-  expr: t.Node,
-  seen: PreviousEnumMembers,
-): number | string | typeof undefined {
-  return evaluate(expr);
+  path: NodePath,
+  prevMembers?: PreviousEnumMembers,
+  seen: Set<t.Identifier> = new Set(),
+): number | string | undefined {
+  return evaluate(path);
 
-  function evaluate(expr: t.Node): number | typeof undefined {
+  function evaluate(path: NodePath): number | string | undefined {
+    const expr = path.node;
     switch (expr.type) {
+      case "MemberExpression":
+        return evaluateRef(path, prevMembers, seen);
       case "StringLiteral":
         return expr.value;
       case "UnaryExpression":
-        return evalUnaryExpression(expr);
+        return evalUnaryExpression(path as NodePath<t.UnaryExpression>);
       case "BinaryExpression":
-        return evalBinaryExpression(expr);
+        return evalBinaryExpression(path as NodePath<t.BinaryExpression>);
       case "NumericLiteral":
         return expr.value;
       case "ParenthesizedExpression":
-        return evaluate(expr.expression);
+        return evaluate(path.get("expression"));
       case "Identifier":
-        return seen.get(expr.name);
-      case "TemplateLiteral":
+        return evaluateRef(path, prevMembers, seen);
+      case "TemplateLiteral": {
         if (expr.quasis.length === 1) {
           return expr.quasis[0].value.cooked;
         }
-      /* falls through */
+
+        const paths = (path as NodePath<t.TemplateLiteral>).get("expressions");
+        const quasis = expr.quasis;
+        let str = "";
+
+        for (let i = 0; i < quasis.length; i++) {
+          str += quasis[i].value.cooked;
+
+          if (i + 1 < quasis.length) {
+            const value = evaluateRef(paths[i], prevMembers, seen);
+            if (value === undefined) return undefined;
+            str += value;
+          }
+        }
+        return str;
+      }
       default:
         return undefined;
     }
   }
 
-  function evalUnaryExpression({
-    argument,
-    operator,
-  }: t.UnaryExpression): number | typeof undefined {
-    const value = evaluate(argument);
+  function evaluateRef(
+    path: NodePath,
+    prevMembers: PreviousEnumMembers,
+    seen: Set<t.Identifier>,
+  ): number | string | undefined {
+    if (path.isMemberExpression()) {
+      const expr = path.node;
+
+      const obj = expr.object;
+      const prop = expr.property;
+      if (
+        !t.isIdentifier(obj) ||
+        (expr.computed ? !t.isStringLiteral(prop) : !t.isIdentifier(prop))
+      ) {
+        return;
+      }
+      const bindingIdentifier = path.scope.getBindingIdentifier(obj.name);
+      const data = ENUMS.get(bindingIdentifier);
+      if (!data) return;
+      // @ts-expect-error checked above
+      return data.get(prop.computed ? prop.value : prop.name);
+    } else if (path.isIdentifier()) {
+      const name = path.node.name;
+
+      let value = prevMembers?.get(name);
+      if (value !== undefined) {
+        return value;
+      }
+
+      if (seen.has(path.node)) return;
+
+      const bindingInitPath = path.resolve(); // It only resolves constant bindings
+      if (bindingInitPath) {
+        seen.add(path.node);
+
+        value = computeConstantValue(bindingInitPath, undefined, seen);
+        prevMembers?.set(name, value);
+        return value;
+      }
+    }
+  }
+
+  function evalUnaryExpression(
+    path: NodePath<t.UnaryExpression>,
+  ): number | string | undefined {
+    const value = evaluate(path.get("argument"));
     if (value === undefined) {
       return undefined;
     }
 
-    switch (operator) {
+    switch (path.node.operator) {
       case "+":
         return value;
       case "-":
@@ -245,17 +313,19 @@ function computeConstantValue(
     }
   }
 
-  function evalBinaryExpression(expr: t.BinaryExpression): number | undefined {
-    const left = evaluate(expr.left);
+  function evalBinaryExpression(
+    path: NodePath<t.BinaryExpression>,
+  ): number | string | undefined {
+    const left = evaluate(path.get("left")) as any;
     if (left === undefined) {
       return undefined;
     }
-    const right = evaluate(expr.right);
+    const right = evaluate(path.get("right")) as any;
     if (right === undefined) {
       return undefined;
     }
 
-    switch (expr.operator) {
+    switch (path.node.operator) {
       case "|":
         return left | right;
       case "&":
diff --git a/packages/babel-plugin-transform-typescript/test/fixtures/enum/mix-references/output.js b/packages/babel-plugin-transform-typescript/test/fixtures/enum/mix-references/output.js
@@ -3,14 +3,14 @@ var Foo;
 (function (Foo) {
   Foo[Foo["a"] = 10] = "a";
   Foo[Foo["b"] = 10] = "b";
-  Foo[Foo["c"] = Foo.b + x] = "c";
+  Foo[Foo["c"] = 20] = "c";
 })(Foo || (Foo = {}));
 var Bar;
 (function (Bar) {
-  Bar[Bar["D"] = Foo.a] = "D";
-  Bar[Bar["E"] = Bar.D] = "E";
+  Bar[Bar["D"] = 10] = "D";
+  Bar[Bar["E"] = 10] = "E";
   Bar[Bar["F"] = Math.E] = "F";
-  Bar[Bar["G"] = Bar.E + Foo.c] = "G";
+  Bar[Bar["G"] = 30] = "G";
 })(Bar || (Bar = {}));
 var Baz;
 (function (Baz) {
diff --git a/packages/babel-plugin-transform-typescript/test/fixtures/enum/outer-references/output.js b/packages/babel-plugin-transform-typescript/test/fixtures/enum/outer-references/output.js
@@ -6,9 +6,9 @@ var socketType;
 })(socketType || (socketType = {}));
 var constants;
 (function (constants) {
-  constants[constants["SOCKET"] = socketType.SOCKET] = "SOCKET";
-  constants[constants["SERVER"] = socketType.SERVER] = "SERVER";
-  constants[constants["IPC"] = socketType.IPC] = "IPC";
-  constants[constants["UV_READABLE"] = 1 + constants["IPC"]] = "UV_READABLE";
-  constants[constants["UV_WRITABLE"] = 1 + constants["UV_READABLE"]] = "UV_WRITABLE";
+  constants[constants["SOCKET"] = 0] = "SOCKET";
+  constants[constants["SERVER"] = 1] = "SERVER";
+  constants[constants["IPC"] = 2] = "IPC";
+  constants[constants["UV_READABLE"] = 3] = "UV_READABLE";
+  constants[constants["UV_WRITABLE"] = 4] = "UV_WRITABLE";
 })(constants || (constants = {}));
diff --git a/packages/babel-plugin-transform-typescript/test/fixtures/enum/ts5.0-const-foldable/input.ts b/packages/babel-plugin-transform-typescript/test/fixtures/enum/ts5.0-const-foldable/input.ts
@@ -0,0 +1,19 @@
+// ts 5.0 feature
+
+const BaseValue = 10;
+const Prefix = "/data";
+const enum Values {
+  First = BaseValue, // 10
+  Second, // 11
+  Third, // 12
+}
+
+const xxx = 100 + Values.First;
+const yyy = xxx;
+
+const enum Routes {
+  Parts = `${Prefix}/parts`, // "/data/parts"
+  Invoices = `${Prefix}/invoices`, // "/data/invoices"
+  x = `${Values.First}/x`,
+  y = `${yyy}/y`,
+}
diff --git a/packages/babel-plugin-transform-typescript/test/fixtures/enum/ts5.0-const-foldable/output.js b/packages/babel-plugin-transform-typescript/test/fixtures/enum/ts5.0-const-foldable/output.js

Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,7 @@ export default function transpileConstEnum(`
`27`	`27`	`);`
`28`	`28`	`}`
`29`	`29`
`30`		`- const entries = translateEnumValues(path, t);`
	`30`	`+ const { enumValues: entries } = translateEnumValues(path, t);`
`31`	`31`
`32`	`32`	`if (isExported) {`
`33`	`33`	`const obj = t.objectExpression(`