Optimize TS enums output (#15467)

Co-authored-by: Huáng Jùnliàng <jlhwung@gmail.com>
Co-authored-by: Nicolò Ribaudo <nicolo.ribaudo@gmail.com>

packages/babel-plugin-transform-typescript/package.json

@@ -17,6 +17,7 @@
     "typescript"
   ],
   "dependencies": {
+    "@babel/helper-annotate-as-pure": "workspace:^",
     "@babel/helper-create-class-features-plugin": "workspace:^",
     "@babel/helper-plugin-utils": "workspace:^",
     "@babel/plugin-syntax-typescript": "workspace:^"

packages/babel-plugin-transform-typescript/src/enum.ts

@@ -1,41 +1,68 @@
 import { template, types as t } from "@babel/core";
 import type { NodePath } from "@babel/traverse";
 import assert from "assert";
+import annotateAsPure from "@babel/helper-annotate-as-pure";
 
 type t = typeof t;
 
 const ENUMS = new WeakMap<t.Identifier, PreviousEnumMembers>();
 
+const buildEnumWrapper = template.expression(
+  `
+    (function (ID) {
+      ASSIGNMENTS;
+      return ID;
+    })(INIT)
+  `,
+);
+
 export default function transpileEnum(
   path: NodePath<t.TSEnumDeclaration>,
   t: t,
 ) {
-  const { node } = path;
+  const { node, parentPath } = path;
 
   if (node.declare) {
     path.remove();
     return;
   }
 
   const name = node.id.name;
-  const { wrapper: fill, data } = enumFill(path, t, node.id);
+  const { fill, data, isPure } = enumFill(path, t, node.id);
 
-  switch (path.parent.type) {
+  switch (parentPath.type) {
     case "BlockStatement":
     case "ExportNamedDeclaration":
     case "Program": {
-      path.insertAfter(fill);
-      if (seen(path.parentPath)) {
-        path.remove();
+      // todo: Consider exclude program with import/export
+      // && !path.parent.body.some(n => t.isImportDeclaration(n) || t.isExportDeclaration(n));
+      const isGlobal = t.isProgram(path.parent);
+      const isSeen = seen(parentPath);
+
+      let init: t.Expression = t.objectExpression([]);
+      if (isSeen || isGlobal) {
+        init = t.logicalExpression("||", t.cloneNode(fill.ID), init);
+      }
+      const enumIIFE = buildEnumWrapper({ ...fill, INIT: init });
+      if (isPure) annotateAsPure(enumIIFE);
+
+      if (isSeen) {
+        const toReplace = parentPath.isExportDeclaration() ? parentPath : path;
+        toReplace.replaceWith(
+          t.expressionStatement(
+            t.assignmentExpression("=", t.cloneNode(node.id), enumIIFE),
+          ),
+        );
       } else {
-        // todo: Consider exclude program with import/export
-        // && !path.parent.body.some(n => t.isImportDeclaration(n) || t.isExportDeclaration(n));
-        const isGlobal = t.isProgram(path.parent);
         path.scope.registerDeclaration(
-          path.replaceWith(makeVar(node.id, t, isGlobal ? "var" : "let"))[0],
+          path.replaceWith(
+            t.variableDeclaration(isGlobal ? "var" : "let", [
+              t.variableDeclarator(node.id, enumIIFE),
+            ]),
+          )[0],
         );
-        ENUMS.set(path.scope.getBindingIdentifier(name), data);
       }
+      ENUMS.set(path.scope.getBindingIdentifier(name), data);
       break;
     }
 
@@ -57,16 +84,6 @@ export default function transpileEnum(
   }
 }
 
-function makeVar(id: t.Identifier, t: t, kind: "var" | "let" | "const") {
-  return t.variableDeclaration(kind, [t.variableDeclarator(id)]);
-}
-
-const buildEnumWrapper = template(`
-  (function (ID) {
-    ASSIGNMENTS;
-  })(ID || (ID = {}));
-`);
-
 const buildStringAssignment = template(`
   ENUM["NAME"] = VALUE;
 `);
@@ -83,7 +100,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 { enumValues: x, data } = translateEnumValues(path, t);
+  const { enumValues: x, data, isPure } = translateEnumValues(path, t);
   const assignments = x.map(([memberName, memberValue]) =>
     buildEnumMember(t.isStringLiteral(memberValue), {
       ENUM: t.cloneNode(id),
@@ -93,11 +110,12 @@ function enumFill(path: NodePath<t.TSEnumDeclaration>, t: t, id: t.Identifier) {
   );
 
   return {
-    wrapper: buildEnumWrapper({
+    fill: {
       ID: t.cloneNode(id),
       ASSIGNMENTS: assignments,
-    }),
-    data: data,
+    },
+    data,
+    isPure,
   };
 }
 
@@ -141,10 +159,11 @@ export function translateEnumValues(path: NodePath<t.TSEnumDeclaration>, t: t) {
   // Start at -1 so the first enum member is its increment, 0.
   let constValue: number | string | undefined = -1;
   let lastName: string;
+  let isPure = true;
 
-  return {
-    data: seen,
-    enumValues: path.get("members").map(memberPath => {
+  const enumValues: Array<[name: string, value: t.Expression]> = 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");
@@ -161,6 +180,8 @@ export function translateEnumValues(path: NodePath<t.TSEnumDeclaration>, t: t) {
             value = t.stringLiteral(constValue);
           }
         } else {
+          isPure &&= initializerPath.isPure();
+
           if (initializerPath.isReferencedIdentifier()) {
             ReferencedIdentifier(initializerPath, {
               t,
@@ -197,7 +218,12 @@ export function translateEnumValues(path: NodePath<t.TSEnumDeclaration>, t: t) {
 
       lastName = name;
       return [name, value];
-    }) as Array<[name: string, value: t.Expression]>,
+    });
+
+  return {
+    isPure,
+    data: seen,
+    enumValues,
   };
 }
 

packages/babel-plugin-transform-typescript/test/fixtures/enum/boolean-value/output.js

@@ -1,5 +1,5 @@
 // Not type-correct code
-var E;
-(function (E) {
+var E = /*#__PURE__*/function (E) {
   E[E["A"] = true] = "A";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/const/output.js

@@ -1,3 +1,4 @@
 // With --isolatedModules, TSC ignores the "const" modifier when compiling enums
-var E;
-(function (E) {})(E || (E = {}));
+var E = /*#__PURE__*/function (E) {
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/constant-folding/output.js

@@ -1,5 +1,4 @@
-var E;
-(function (E) {
+var E = /*#__PURE__*/function (E) {
   E[E["a"] = 0] = "a";
   E[E["b"] = 3] = "b";
   E[E["c"] = 1] = "c";
@@ -18,4 +17,5 @@ var E;
   E[E["p"] = -3] = "p";
   E[E["q"] = 5.4] = "q";
   E[E["r"] = 6.4] = "r";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/export/output.mjs

@@ -1,7 +1,8 @@
-export let E;
-(function (E) {
+export let E = /*#__PURE__*/function (E) {
   E[E["A"] = 1] = "A";
-})(E || (E = {}));
-(function (E) {
+  return E;
+}({});
+E = /*#__PURE__*/function (E) {
   E[E["B"] = 2] = "B";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/inferred/output.js

@@ -1,5 +1,5 @@
-var E;
-(function (E) {
+var E = /*#__PURE__*/function (E) {
   E[E["x"] = 0] = "x";
   E[E["y"] = 1] = "y";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/inner-references/output.js

@@ -1,5 +1,5 @@
-var E;
-(function (E) {
+var E = /*#__PURE__*/function (E) {
   E[E["a"] = 10] = "a";
   E[E["b"] = 10] = "b";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/mix-references/output.js

@@ -1,25 +1,24 @@
 var x = 10;
-var Foo;
-(function (Foo) {
+var Foo = /*#__PURE__*/function (Foo) {
   Foo[Foo["a"] = 10] = "a";
   Foo[Foo["b"] = 10] = "b";
   Foo[Foo["c"] = 20] = "c";
-})(Foo || (Foo = {}));
-var Bar;
-(function (Bar) {
+  return Foo;
+}(Foo || {});
+var Bar = function (Bar) {
   Bar[Bar["D"] = 10] = "D";
   Bar[Bar["E"] = 10] = "E";
   Bar[Bar["F"] = Math.E] = "F";
   Bar[Bar["G"] = 30] = "G";
-})(Bar || (Bar = {}));
-var Baz;
-(function (Baz) {
+  return Bar;
+}(Bar || {});
+var Baz = function (Baz) {
   Baz[Baz["a"] = 0] = "a";
   Baz[Baz["b"] = 1] = "b";
   Baz[Baz["x"] = Baz.a.toString()] = "x";
-})(Baz || (Baz = {}));
-var A; // a refers to A.a
-(function (A) {
+  return Baz;
+}(Baz || {});
+var A = function (A) {
   A[A["a"] = 0] = "a";
   A[A["b"] = (() => {
     let a = 1;
@@ -28,4 +27,5 @@ var A; // a refers to A.a
   A[A["c"] = (() => {
     return A.a + 2;
   })()] = "c";
-})(A || (A = {}));
+  return A;
+}(A || {}); // a refers to A.a

packages/babel-plugin-transform-typescript/test/fixtures/enum/non-foldable-constant/output.js

@@ -1,5 +1,5 @@
-var E;
-(function (E) {
+var E = function (E) {
   E[E["a"] = Math.sin(1)] = "a";
   E[E["b"] = 1 + E["a"]] = "b";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/non-scoped/output.js

@@ -1,8 +1,9 @@
-var E;
-(function (E) {
+var E = /*#__PURE__*/function (E) {
   E[E["x"] = 1] = "x";
   E[E["y"] = 2] = "y";
-})(E || (E = {}));
-(function (E) {
+  return E;
+}(E || {});
+E = /*#__PURE__*/function (E) {
   E[E["z"] = 3] = "z";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/outer-references/output.js

@@ -1,14 +1,14 @@
-var socketType;
-(function (socketType) {
+var socketType = /*#__PURE__*/function (socketType) {
   socketType[socketType["SOCKET"] = 0] = "SOCKET";
   socketType[socketType["SERVER"] = 1] = "SERVER";
   socketType[socketType["IPC"] = 2] = "IPC";
-})(socketType || (socketType = {}));
-var constants;
-(function (constants) {
+  return socketType;
+}(socketType || {});
+var constants = /*#__PURE__*/function (constants) {
   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 = {}));
+  return constants;
+}(constants || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/scoped/output.js

@@ -1,5 +1,6 @@
 {
   // Uses 'let' within a scope
-  let E;
-  (function (E) {})(E || (E = {}));
+  let E = /*#__PURE__*/function (E) {
+    return E;
+  }({});
 }

packages/babel-plugin-transform-typescript/test/fixtures/enum/string-value-template/output.js

@@ -1,4 +1,4 @@
-var E;
-(function (E) {
+var E = /*#__PURE__*/function (E) {
   E["A"] = "Hey";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/string-value/output.js

@@ -1,7 +1,7 @@
-var E;
-(function (E) {
+var E = /*#__PURE__*/function (E) {
   E[E["A"] = 0] = "A";
   E["B"] = "";
   E[E["A2"] = 0] = "A2";
   E["B2"] = "";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/string-values-computed/output.js

@@ -1,4 +1,4 @@
-var E;
-(function (E) {
+var E = /*#__PURE__*/function (E) {
   E["A"] = "HALLOWERLD";
-})(E || (E = {}));
+  return E;
+}(E || {});

packages/babel-plugin-transform-typescript/test/fixtures/enum/ts5.0-const-foldable/output.js

@@ -2,18 +2,18 @@
 
 const BaseValue = 10;
 const Prefix = "/data";
-var Values; // 12
-(function (Values) {
+var Values = /*#__PURE__*/function (Values) {
   Values[Values["First"] = 10] = "First";
   Values[Values["Second"] = 11] = "Second";
   Values[Values["Third"] = 12] = "Third";
-})(Values || (Values = {}));
+  return Values;
+}(Values || {}); // 12
 const xxx = 100 + Values.First;
 const yyy = xxx;
-var Routes;
-(function (Routes) {
+var Routes = /*#__PURE__*/function (Routes) {
   Routes["Parts"] = "/data/parts";
   Routes["Invoices"] = "/data/invoices";
   Routes["x"] = "10/x";
   Routes["y"] = "110/y";
-})(Routes || (Routes = {}));
+  return Routes;
+}(Routes || {});