Simplify transform-react-jsx-source code (#15046)

babel · Oct 17, 2022 · ddaee90 · ddaee90
1 parent 3526b79
commit ddaee90
Showing 1 changed file with 45 additions and 65 deletions.
diff --git a/packages/babel-plugin-transform-react-jsx-source/src/index.ts b/packages/babel-plugin-transform-react-jsx-source/src/index.ts
@@ -13,12 +13,16 @@
  * <sometag __source={{fileName: __jsxFileName, lineNumber: 10, columnNumber: 1}}/>
  */
 import { declare } from "@babel/helper-plugin-utils";
-import { type PluginPass, types as t } from "@babel/core";
-import type { Visitor } from "@babel/traverse";
+import { types as t, template } from "@babel/core";
 
 const TRACE_ID = "__source";
 const FILE_NAME_VAR = "_jsxFileName";
 
+const createNodeFromNullish = <T, N extends t.Node>(
+  val: T | null,
+  fn: (val: T) => N,
+): N | t.NullLiteral => (val == null ? t.nullLiteral() : fn(val));
+
 type State = {
   fileNameIdentifier: t.Identifier;
 };
@@ -27,81 +31,57 @@ export default declare<State>(api => {
 
   function makeTrace(
     fileNameIdentifier: t.Identifier,
-    lineNumber: number,
-    column0Based: number,
+    { line, column }: { line: number; column: number },
   ) {
-    const fileLineLiteral =
-      lineNumber != null ? t.numericLiteral(lineNumber) : t.nullLiteral();
-    const fileColumnLiteral =
-      column0Based != null
-        ? t.numericLiteral(column0Based + 1)
-        : t.nullLiteral();
-    const fileNameProperty = t.objectProperty(
-      t.identifier("fileName"),
-      fileNameIdentifier,
-    );
-    const lineNumberProperty = t.objectProperty(
-      t.identifier("lineNumber"),
-      fileLineLiteral,
-    );
-    const columnNumberProperty = t.objectProperty(
-      t.identifier("columnNumber"),
-      fileColumnLiteral,
+    const fileLineLiteral = createNodeFromNullish(line, t.numericLiteral);
+    const fileColumnLiteral = createNodeFromNullish(column, c =>
+      // c + 1 to make it 1-based instead of 0-based.
+      t.numericLiteral(c + 1),
     );
-    return t.objectExpression([
-      fileNameProperty,
-      lineNumberProperty,
-      columnNumberProperty,
-    ]);
+
+    return template.expression.ast`{
+      fileName: ${fileNameIdentifier},
+      lineNumber: ${fileLineLiteral},
+      columnNumber: ${fileColumnLiteral},
+    }`;
   }
 
-  const visitor: Visitor<State & PluginPass> = {
-    JSXOpeningElement(path, state) {
-      const id = t.jsxIdentifier(TRACE_ID);
-      const location = (path.container as t.JSXElement).openingElement.loc;
-      if (!location) {
-        // the element was generated and doesn't have location information
-        return;
-      }
+  const isSourceAttr = (attr: t.Node) =>
+    t.isJSXAttribute(attr) && attr.name.name === TRACE_ID;
 
-      const attributes = (path.container as t.JSXElement).openingElement
-        .attributes;
-      for (let i = 0; i < attributes.length; i++) {
-        // @ts-expect-error .name is not defined in JSXSpreadElement
-        const name = attributes[i].name as t.JSXAttribute["name"] | void;
-        // @ts-expect-error TS can not narrow down optional chain
-        if (name?.name === TRACE_ID) {
-          // The __source attribute already exists
+  return {
+    name: "transform-react-jsx-source",
+    visitor: {
+      JSXOpeningElement(path, state) {
+        const { node } = path;
+        if (
+          // the element was generated and doesn't have location information
+          !node.loc ||
+          // Already has __source
+          path.node.attributes.some(isSourceAttr)
+        ) {
           return;
         }
-      }
 
-      if (!state.fileNameIdentifier) {
-        const fileName = state.filename || "";
+        if (!state.fileNameIdentifier) {
+          const fileNameId = path.scope.generateUidIdentifier(FILE_NAME_VAR);
+          state.fileNameIdentifier = fileNameId;
 
-        const fileNameIdentifier =
-          path.scope.generateUidIdentifier(FILE_NAME_VAR);
-        const scope = path.hub.getScope();
-        if (scope) {
-          scope.push({
-            id: fileNameIdentifier,
-            init: t.stringLiteral(fileName),
+          path.scope.getProgramParent().push({
+            id: fileNameId,
+            init: t.stringLiteral(state.filename || ""),
           });
         }
-        state.fileNameIdentifier = fileNameIdentifier;
-      }
 
-      const trace = makeTrace(
-        t.cloneNode(state.fileNameIdentifier),
-        location.start.line,
-        location.start.column,
-      );
-      attributes.push(t.jsxAttribute(id, t.jsxExpressionContainer(trace)));
+        node.attributes.push(
+          t.jsxAttribute(
+            t.jsxIdentifier(TRACE_ID),
+            t.jsxExpressionContainer(
+              makeTrace(t.cloneNode(state.fileNameIdentifier), node.loc.start),
+            ),
+          ),
+        );
+      },
     },
   };
-
-  return {
-    name: "transform-react-jsx-source",
-    visitor,
-  };
 });