Support wrapGenerator.isGeneratorFunction

lib/visit.js

@@ -38,6 +38,7 @@ function visitNode(node) {
   // TODO Ensure these identifiers are named uniquely.
   var contextId = b.identifier("$ctx");
   var argsId = b.identifier("$args");
+  var wrapGeneratorId = b.identifier("wrapGenerator");
   var shouldAliasArguments = renameArguments(node, argsId);
   var vars = hoist(node);
 
@@ -59,13 +60,39 @@ function visitNode(node) {
   }
 
   outerBody.push(b.returnStatement(
-    b.callExpression(b.identifier("wrapGenerator"), [
+    b.callExpression(wrapGeneratorId, [
       emitter.getContextFunction(),
       b.thisExpression()
     ])
   ));
 
   node.body = b.blockStatement(outerBody);
+
+  var markMethod = b.memberExpression(
+    wrapGeneratorId,
+    b.identifier("mark"),
+    false
+  );
+
+  if (n.FunctionDeclaration.check(node)) {
+    var path = this.parent;
+
+    while (path && !n.BlockStatement.check(path.value)) {
+      path = path.parent;
+    }
+
+    if (path) {
+      var firstStmtPath = path.get("body", 0);
+      firstStmtPath.replace(
+        b.expressionStatement(b.callExpression(markMethod, [node.id])),
+        firstStmtPath.value
+      );
+    }
+
+  } else {
+    n.FunctionExpression.assert(node);
+    this.replace(b.callExpression(markMethod, [node]));
+  }
 }
 
 function renameArguments(func, argsId) {

runtime/dev.js

@@ -32,6 +32,20 @@
   // breaking out of the dispatch switch statement.
   var ContinueSentinel = {};
 
+  // Dummy constructor that we use as the .constructor property for
+  // functions that return Generator objects.
+  function GeneratorFunction() {}
+
+  wrapGenerator.mark = function(genFun) {
+    genFun.constructor = GeneratorFunction;
+    return genFun;
+  };
+
+  wrapGenerator.isGeneratorFunction = function(genFun) {
+    var ctor = genFun && genFun.constructor;
+    return ctor ? "GeneratorFunction" === ctor.name : false;
+  };
+
   function Generator(innerFn, self) {
     var generator = this;
     var context = new Context();

test/tests.es6.js

@@ -963,3 +963,43 @@ describe("yield* expression results", function () {
     assert.strictEqual(pumpNumber(foo()), 10);
   });
 });
+
+describe("isGeneratorFunction", function() {
+  it("should work for function declarations", function() {
+    // Do the assertions up here to make sure the generator function is
+    // marked at the beginning of the block the function is declared in.
+    assert.strictEqual(
+      wrapGenerator.isGeneratorFunction(genFun),
+      true
+    );
+
+    assert.strictEqual(
+      wrapGenerator.isGeneratorFunction(normalFun),
+      false
+    );
+
+    function normalFun() {
+      return 0;
+    }
+
+    function *genFun() {
+      yield 0;
+    }
+  });
+
+  it("should work for function expressions", function() {
+    assert.strictEqual(
+      wrapGenerator.isGeneratorFunction(function *genFun() {
+        yield 0;
+      }),
+      true
+    );
+
+    assert.strictEqual(
+      wrapGenerator.isGeneratorFunction(function normalFun() {
+        return 0;
+      }),
+      false
+    );
+  });
+});