refactor mark declaration, add __PURE__ comments (#286)

packages/regenerator-transform/src/visit.js

@@ -153,6 +153,7 @@ exports.visitor = {
 
       if (wasGeneratorFunction && t.isExpression(node)) {
         util.replaceWithOrRemove(path, t.callExpression(util.runtimeProperty("mark"), [node]))
+        path.addComment("leading", "#__PURE__");
       }
 
       // Generators are processed in 'exit' handlers so that regenerator only has to run on
@@ -188,18 +189,34 @@ function getOuterFnExpr(funPath) {
     }
 
     let markDecl = getRuntimeMarkDecl(pp);
-    let markedArray = markDecl.declarations[0].id;
-    let funDeclIdArray = markDecl.declarations[0].init.callee.object;
-    t.assertArrayExpression(funDeclIdArray);
-
-    let index = funDeclIdArray.elements.length;
-    funDeclIdArray.elements.push(node.id);
-
-    return t.memberExpression(
-      markedArray,
-      t.numericLiteral(index),
-      true
-    );
+    let declarations = markDecl.declarations;
+    let index = declarations.length;
+
+    // this can change from pass to pass if something else is altering our body
+    // so we have to look it up every time :/
+    let declarationsIndex = pp.node.body.findIndex((p) => {
+      return t.isVariableDeclaration(p) && p.declarations === declarations;
+    });
+    assert(!isNaN(declarationsIndex));
+
+    // get a new unique id for our marked variable
+    let markedId = pp.scope.generateUidIdentifier("marked");
+
+    // now push our new declaration into our list
+    declarations.push(t.variableDeclarator(
+      markedId,
+      t.callExpression(util.runtimeProperty("mark"), [node.id])
+    ));
+
+    // assemble the path to our mark declarations
+    let bodyPathName = 'body.' + String(declarationsIndex) + '.declarations.' + String(index) + '.init';
+    let bodyPath = pp.get(bodyPathName);
+
+    // mark our call to util.runtimeProperty("mark") as pure to enable tree-shaking
+    bodyPath.addComment("leading", "#__PURE__");
+
+    // and return our new id
+    return markedId;
   }
 
   return node.id;
@@ -214,20 +231,7 @@ function getRuntimeMarkDecl(blockPath) {
     return info.decl;
   }
 
-  info.decl = t.variableDeclaration("var", [
-    t.variableDeclarator(
-      blockPath.scope.generateUidIdentifier("marked"),
-      t.callExpression(
-        t.memberExpression(
-          t.arrayExpression([]),
-          t.identifier("map"),
-          false
-        ),
-        [util.runtimeProperty("mark")]
-      )
-    )
-  ]);
-
+  info.decl = t.variableDeclaration("var", [])
   blockPath.unshiftContainer("body", info.decl);
 
   return info.decl;

test/tests.transform.js

@@ -47,3 +47,57 @@ describe("_blockHoist nodes", function() {
     assert.deepEqual(names, ["hoistMe", "doNotHoistMe", "oyez"]);
   });
 });
+
+describe("marked", function() {
+  it("should work with a single function", function() {
+    var ast = recast.parse('function* foo(){};', {
+        parser: require("babylon")
+    });
+    const code = recast.print(transform(ast)).code;
+    assert.notStrictEqual(
+      code.indexOf(
+        'var _marked = regeneratorRuntime.mark(foo);'
+      ), -1
+    );
+  });
+  it("should work with multiple functions", function() {
+    var ast = recast.parse([
+      'function* foo(){};',
+      'function* bar() {};'
+    ].join("\n"), {
+        parser: require("babylon")
+    });
+    const code = recast.print(transform(ast)).code;
+    assert.notStrictEqual(
+      code.indexOf(
+        'var _marked = regeneratorRuntime.mark(foo), _marked2 = regeneratorRuntime.mark(bar);'
+      ), -1
+    );
+  });
+});
+
+// TODO - how do i get recast to return comments?
+describe("pure", function() {
+  xit("should work with a function expression", function() {
+    var ast = recast.parse('var a = function* foo(){};', {
+        parser: require("babylon")
+    });
+    const code = recast.print(transform(ast)).code;
+    assert.notStrictEqual(
+      code.indexOf(
+        'var a = /*#__PURE__*/regeneratorRuntime.mark(function foo() {'
+      ), -1
+    );
+  });
+  xit("should work with a function declaration", function() {
+    var ast = recast.parse('function* foo(){};', {
+      parser: require("babylon")
+    });
+    const code = recast.print(transform(ast)).code;
+    assert.notStrictEqual(
+      code.indexOf(
+        'var a = /*#__PURE__*/Object.defineProperty( /*#__PURE__*/regeneratorRuntime.mark'
+      ), -1
+    );
+  });
+});