refactor: improved transformation logic

packages/babel-plugin-transform-optional-chaining/src/index.js

@@ -1,61 +1,44 @@
-import createTemplate from "babel-template";
-import traverse from "babel-traverse";
-
-const nullOrUndefinedCheck = createTemplate(`
-  typeof CHECK !== "undefined" && CHECK !== null
-    ? NEXT
-    : null
-`);
-
-function isNodeOptional(node) {
-  return node.optional === true;
-}
-
-const nullOrUndefinedCheckVisitor = {
-  noScope: true,
-
-  Identifier(path, replacements) {
-    if (path.node.name in replacements) {
-      path.replaceInline(replacements[path.node.name]);
-    }
-  },
-};
-
-function createCheck(node, object) {
-
-  const template = nullOrUndefinedCheck();
-
-  traverse(template, nullOrUndefinedCheckVisitor, null, {
-    CHECK: object,
-    NEXT: node,
-  });
-
-  return template;
-}
-
 export default function ({ types: t }) {
 
+  function createCondition(ref, access, nextProperty) {
+
+    return t.conditionalExpression(
+      createCheckAgainstNull(
+        t.AssignmentExpression("=", ref, access)
+      ),
+
+      t.memberExpression(ref, nextProperty),
+      t.NullLiteral()
+    );
+  }
+
+  function createCheckAgainstNull(left) {
+    return t.BinaryExpression("!=", left, t.NullLiteral());
+  }
+
+  function isNodeOptional(node) {
+    return node.optional === true;
+  }
+
   return {
     visitor: {
-      MemberExpression(path) {
-
-        if (isNodeOptional(path.node)) {
-          let { object } = path.node;
-
-          do {
-            object = createCheck(
-              object,
-              object.object
-            );
-
-          } while (t.isMemberExpression(object) && isNodeOptional(object));
-
-          path.replaceWith(
-            createCheck(path.node, object)
-          );
-
-          path.stop();
+      MemberExpression(path, state) {
+        if (!isNodeOptional(path.node)) {
+          return;
         }
+
+        const { object, property } = path.node;
+
+        if (!state.optionalTemp) {
+          const id = path.scope.generateUidIdentifier();
+
+          state.optionalTemp = id;
+          path.scope.push({ id });
+        }
+
+        path.replaceWith(
+          createCondition(state.optionalTemp, object, property)
+        );
       },
     },
   };