Fix TypeScript Enum self-references

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

@@ -1,9 +1,14 @@
-import assert from "assert";
 import { template } from "@babel/core";
-import type * as t from "@babel/types";
 import type { NodePath } from "@babel/traverse";
+import type * as t from "@babel/types";
+import assert from "assert";
 
-export default function transpileEnum(path, t) {
+type t = typeof t;
+
+export default function transpileEnum(
+  path: NodePath<t.TSEnumDeclaration>,
+  t: t,
+) {
   const { node } = path;
 
   if (node.declare) {
@@ -48,7 +53,7 @@ export default function transpileEnum(path, t) {
   }
 }
 
-function makeVar(id, t, kind) {
+function makeVar(id: t.Identifier, t: t, kind: "var" | "let" | "const") {
   return t.variableDeclaration(kind, [t.variableDeclarator(id)]);
 }
 
@@ -66,14 +71,14 @@ const buildNumericAssignment = template(`
   ENUM[ENUM["NAME"] = VALUE] = "NAME";
 `);
 
-const buildEnumMember = (isString, options) =>
+const buildEnumMember = (isString: boolean, options: Record<string, unknown>) =>
   (isString ? buildStringAssignment : buildNumericAssignment)(options);
 
 /**
  * Generates the statement that fills in the variable declared by the enum.
  * `(function (E) { ... assignments ... })(E || (E = {}));`
  */
-function enumFill(path, t, id) {
+function enumFill(path: NodePath<t.TSEnumDeclaration>, t: t, id: t.Identifier) {
   const x = translateEnumValues(path, t);
   const assignments = x.map(([memberName, memberValue]) =>
     buildEnumMember(t.isStringLiteral(memberValue), {
@@ -100,16 +105,41 @@ function enumFill(path, t, id) {
  */
 type PreviousEnumMembers = Map<string, number | string>;
 
+type EnumSelfReferenceVisitorState = {
+  seen: PreviousEnumMembers;
+  path: NodePath<t.TSEnumDeclaration>;
+  t: t;
+};
+
+function ReferencedIdentifier(
+  expr: NodePath<t.Identifier>,
+  state: EnumSelfReferenceVisitorState,
+) {
+  const { seen, path, t } = state;
+  const name = expr.node.name;
+  if (seen.has(name) && !expr.scope.hasOwnBinding(name)) {
+    expr.replaceWith(
+      t.memberExpression(t.cloneNode(path.node.id), t.cloneNode(expr.node)),
+    );
+    expr.skip();
+  }
+}
+
+const enumSelfReferenceVisitor = {
+  ReferencedIdentifier,
+};
+
 export function translateEnumValues(
   path: NodePath<t.TSEnumDeclaration>,
-  t: typeof import("@babel/types"),
+  t: t,
 ): Array<[name: string, value: t.Expression]> {
   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.node.members.map(member => {
+  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;
@@ -124,7 +154,20 @@ export function translateEnumValues(
           value = t.stringLiteral(constValue);
         }
       } else {
-        value = initializer;
+        const initializerPath = memberPath.get("initializer");
+
+        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;
@@ -140,6 +183,7 @@ export function translateEnumValues(
         true,
       );
       value = t.binaryExpression("+", t.numericLiteral(1), lastRef);
+      seen.set(name, undefined);
     }
 
     lastName = name;

packages/babel-plugin-transform-typescript/test/fixtures/enum/mix-references/input.ts

@@ -0,0 +1,27 @@
+var x = 10;
+
+enum Foo {
+    a = 10,
+    b = a,
+    c = b + x,
+}
+
+enum Bar {
+    D = Foo.a,
+    E = D,
+    F = Math.E,
+    G = E + Foo.c,
+}
+
+enum Baz {
+    a = 0,
+    b = 1,
+    // @ts-ignore
+    x = a.toString(),
+}
+
+enum A {
+    a = 0,
+    b = (() => { let a = 1; return a + 1 })(), // a is shadowed
+    c = (() => { return a + 2 })(), // a refers to A.a
+}

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

@@ -0,0 +1,38 @@
+var x = 10;
+var Foo;
+
+(function (Foo) {
+  Foo[Foo["a"] = 10] = "a";
+  Foo[Foo["b"] = 10] = "b";
+  Foo[Foo["c"] = Foo.b + x] = "c";
+})(Foo || (Foo = {}));
+
+var Bar;
+
+(function (Bar) {
+  Bar[Bar["D"] = Foo.a] = "D";
+  Bar[Bar["E"] = Bar.D] = "E";
+  Bar[Bar["F"] = Math.E] = "F";
+  Bar[Bar["G"] = Bar.E + Foo.c] = "G";
+})(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;
+
+(function (A) {
+  A[A["a"] = 0] = "a";
+  A[A["b"] = (() => {
+    let a = 1;
+    return a + 1;
+  })()] = "b";
+  A[A["c"] = (() => {
+    return A.a + 2;
+  })()] = "c";
+})(A || (A = {}));