add support for computed mutators in babel-plugin-transform-es2015-computed-properties

packages/babel-plugin-transform-es2015-computed-properties/src/index.js

@@ -1,4 +1,9 @@
-export default function ({ types: t }) {
+export default function ({ types: t, template }) {
+  let buildMutatorMapAssign = template(`
+    MUTATOR_MAP_REF[KEY] = MUTATOR_MAP_REF[KEY] || {};
+    MUTATOR_MAP_REF[KEY].KIND = VALUE;
+  `);
+
   function getValue(prop) {
     if (t.isObjectProperty(prop)) {
       return prop.value;
@@ -8,31 +13,60 @@ export default function ({ types: t }) {
   }
 
   function pushAssign(objId, prop, body) {
-    body.push(t.expressionStatement(
-      t.assignmentExpression(
-        "=",
-        t.memberExpression(objId, prop.key, prop.computed || t.isLiteral(prop.key)),
-        getValue(prop)
-      )
-    ));
-  }
-
-  function loose(objId, body, computedProps: Array<Object>) {
-    for (let prop of computedProps) {
-      pushAssign(objId, prop, body);
+    if (prop.kind === "get" && prop.kind === "set") {
+      pushMutatorDefine(objId, prop, body);
+    } else {
+      body.push(t.expressionStatement(
+        t.assignmentExpression(
+          "=",
+          t.memberExpression(objId, prop.key, prop.computed || t.isLiteral(prop.key)),
+          getValue(prop)
+        )
+      ));
     }
   }
 
-  function spec(objId, body, computedProps: Array<Object>, initPropExpression, state) {
+  function pushMutatorDefine({ objId, body, getMutatorId, scope }, prop) {
+    let key = prop.key;
+
+    let maybeMemoise = scope.maybeGenerateMemoised(prop.key);
+    if (maybeMemoise) {
+      body.push(t.expressionStatement(t.assignmentExpression("=", maybeMemoise, key)));
+      key = maybeMemoise;
+    }
+
+    body.push(...buildMutatorMapAssign({
+      MUTATOR_MAP_REF: getMutatorId(),
+      KEY: key,
+      VALUE: getValue(prop),
+      KIND: t.identifier(prop.kind)
+    }));
+  }
+
+  function loose(info) {
+    for (let prop of info.computedProps) {
+      if (prop.kind === "get" || prop.kind === "set") {
+        pushMutatorDefine(info, prop);
+      } else {
+        pushAssign(info.objId, prop, info.body);
+      }
+    }
+  }
+
+  function spec(info) {
+    let { objId, body, computedProps, state } = info;
+
     for (let prop of computedProps) {
       let key = t.toComputedKey(prop);
 
-      if (t.isStringLiteral(key, { value: "__proto__" })) {
+      if (prop.kind === "get" || prop.kind === "set") {
+        pushMutatorDefine(info, prop);
+      } else if (t.isStringLiteral(key, { value: "__proto__" })) {
         pushAssign(objId, prop, body);
       } else {
         if (computedProps.length === 1) {
           return t.callExpression(state.addHelper("defineProperty"), [
-            initPropExpression,
+            info.initPropExpression,
             key,
             getValue(prop)
           ]);
@@ -53,10 +87,10 @@ export default function ({ types: t }) {
     visitor: {
       ObjectExpression: {
         exit(path, state) {
-          let { node, parent } = path;
+          let { node, parent, scope } = path;
           let hasComputed = false;
           for (let prop of (node.properties: Array<Object>)) {
-            hasComputed = prop.kind !== "get" && prop.kind !== "set" && prop.computed === true;
+            hasComputed = prop.computed === true;
             if (hasComputed) break;
           }
           if (!hasComputed) return;
@@ -73,14 +107,14 @@ export default function ({ types: t }) {
               foundComputed = true;
             }
 
-            if (foundComputed && prop.kind !== "get" && prop.kind !== "set") {
+            if (foundComputed) {
               computedProps.push(prop);
             } else {
               initProps.push(prop);
             }
           }
 
-          let objId = path.scope.generateUidIdentifierBasedOnNode(parent);
+          let objId = scope.generateUidIdentifierBasedOnNode(parent);
           let initPropExpression = t.objectExpression(initProps);
           let body = [];
 
@@ -91,7 +125,37 @@ export default function ({ types: t }) {
           let callback = spec;
           if (state.opts.loose) callback = loose;
 
-          let single = callback(objId, body, computedProps, initPropExpression, state);
+          let mutatorRef;
+
+          let getMutatorId = function () {
+            if (!mutatorRef) {
+              mutatorRef = scope.generateUidIdentifier("mutatorMap");
+
+              body.push(t.variableDeclaration("var", [
+                t.variableDeclarator(mutatorRef, t.objectExpression([]))
+              ]));
+            }
+
+            return mutatorRef;
+          };
+
+          let single = callback({
+            scope,
+            objId,
+            body,
+            computedProps,
+            initPropExpression,
+            getMutatorId,
+            state,
+          });
+
+          if (mutatorRef) {
+            body.push(t.expressionStatement(t.callExpression(
+              state.addHelper("defineEnumerableProperties"),
+              [objId, mutatorRef]
+            )));
+          }
+
           if (single) {
             path.replaceWith(single);
           } else {