Unify logic for pure and non-pure enums

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,32 +1,26 @@
 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.statement(`
-  (function (ID) {
-    ASSIGNMENTS;
-  })(ID || (ID = {}));
-`);
-
-const buildPureEnumWrapper = template.expression(
+const buildEnumWrapper = template.expression(
   `
-  /*#__PURE__*/(function (ID) {
-    ASSIGNMENTS;
-    return ID;
-  })(ID || {})
-`,
-  { preserveComments: true },
+    (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();
@@ -36,46 +30,37 @@ export default function transpileEnum(
   const name = node.id.name;
   const { fill, data, isPure } = enumFill(path, t, node.id);
 
-  switch (path.parent.type) {
+  switch (parentPath.type) {
     case "BlockStatement":
     case "ExportNamedDeclaration":
     case "Program": {
       // 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);
 
-      if (isPure && isGlobal) {
-        if (seen(path.parentPath)) {
-          path.replaceWith(
-            t.assignmentExpression(
-              "=",
-              t.cloneNode(node.id),
-              buildPureEnumWrapper(fill),
-            ),
-          );
-        } else {
-          path.scope.registerDeclaration(
-            path.replaceWith(
-              t.variableDeclaration("var", [
-                t.variableDeclarator(node.id, buildPureEnumWrapper(fill)),
-              ]),
-            )[0],
-          );
-        }
+      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 {
-        path.insertAfter(buildEnumWrapper(fill));
-
-        if (seen(path.parentPath)) {
-          path.remove();
-        } else {
-          path.scope.registerDeclaration(
-            path.replaceWith(
-              t.variableDeclaration(isGlobal ? "var" : "let", [
-                t.variableDeclarator(node.id),
-              ]),
-            )[0],
-          );
-        }
+        path.scope.registerDeclaration(
+          path.replaceWith(
+            t.variableDeclaration(isGlobal ? "var" : "let", [
+              t.variableDeclarator(node.id, enumIIFE),
+            ]),
+          )[0],
+        );
       }
       ENUMS.set(path.scope.getBindingIdentifier(name), data);
       break;

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/mix-references/output.js

@@ -5,21 +5,20 @@ var Foo = /*#__PURE__*/function (Foo) {
   Foo[Foo["c"] = 20] = "c";
   return Foo;
 }(Foo || {});
-var Bar;
-(function (Bar) {
+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/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/namespace/nested/output.mjs

@@ -14,20 +14,20 @@ class A {}
   _A.D = D;
   (function (_D) {
     const C = 5;
-    let H;
-    (function (H) {
+    let H = /*#__PURE__*/function (H) {
       H[H["I"] = 11] = "I";
       H[H["J"] = 13] = "J";
       H[H["K"] = 17] = "K";
-    })(H || (H = {}));
+      return H;
+    }({});
     _D.H = H;
   })(D || (D = _A.D || (_A.D = {})));
   class F {}
   (function (_F) {})(F || (F = {}));
   let G;
   (function (_G) {})(G || (G = {}));
-  let L;
-  (function (L) {
+  let L = /*#__PURE__*/function (L) {
     L[L["M"] = 19] = "M";
-  })(L || (L = {}));
+    return L;
+  }({});
 })(A || (A = {}));

packages/babel-plugin-transform-typescript/test/fixtures/namespace/same-name/output.mjs

@@ -12,8 +12,9 @@ let N;
     _N8._N5 = _N5;
   })(N || (N = _N2.N || (_N2.N = {})));
   (function (_N9) {
-    let _N;
-    (function (_N) {})(_N || (_N = {}));
+    let _N = /*#__PURE__*/function (_N) {
+      return _N;
+    }({});
     _N9._N = _N;
   })(N || (N = _N2.N || (_N2.N = {})));
 })(N || (N = {}));

yarn.lock

@@ -3303,6 +3303,7 @@ __metadata:
   resolution: "@babel/plugin-transform-typescript@workspace:packages/babel-plugin-transform-typescript"
   dependencies:
     "@babel/core": "workspace:^"
+    "@babel/helper-annotate-as-pure": "workspace:^"
     "@babel/helper-create-class-features-plugin": "workspace:^"
     "@babel/helper-plugin-test-runner": "workspace:^"
     "@babel/helper-plugin-utils": "workspace:^"