Implement support for async generator functions and for-await statements

packages/babel-generator/src/generators/statements.js

@@ -81,6 +81,10 @@ export function WhileStatement(node: Object) {
 let buildForXStatement = function (op) {
   return function (node: Object) {
     this.keyword("for");
+    if (op === "await") {
+      this.push("await ");
+      op = "of";
+    }
     this.push("(");
     this.print(node.left, node);
     this.push(` ${op} `);
@@ -92,6 +96,7 @@ let buildForXStatement = function (op) {
 
 export let ForInStatement = buildForXStatement("in");
 export let ForOfStatement = buildForXStatement("of");
+export let ForAwaitStatement = buildForXStatement("await");
 
 export function DoWhileStatement(node: Object) {
   this.push("do ");

packages/babel-helper-remap-async-to-generator/src/for-await.js

@@ -0,0 +1,106 @@
+import * as t from "babel-types";
+import template from "babel-template";
+import traverse from "babel-traverse";
+
+let buildForAwait = template(`
+  function* wrapper() {
+    var ITERATOR_COMPLETION = true;
+    var ITERATOR_HAD_ERROR_KEY = false;
+    var ITERATOR_ERROR_KEY = undefined;
+    try {
+      for (
+        var ITERATOR_KEY = GET_ITERATOR(OBJECT), STEP_KEY, STEP_VALUE;
+        (
+          STEP_KEY = yield AWAIT(ITERATOR_KEY.next()),
+          ITERATOR_COMPLETION = STEP_KEY.done,
+          STEP_VALUE = yield AWAIT(STEP_KEY.value),
+          !ITERATOR_COMPLETION
+        );
+        ITERATOR_COMPLETION = true) {
+      }
+    } catch (err) {
+      ITERATOR_HAD_ERROR_KEY = true;
+      ITERATOR_ERROR_KEY = err;
+    } finally {
+      try {
+        if (!ITERATOR_COMPLETION && ITERATOR_KEY.return) {
+          yield AWAIT(ITERATOR_KEY.return());
+        }
+      } finally {
+        if (ITERATOR_HAD_ERROR_KEY) {
+          throw ITERATOR_ERROR_KEY;
+        }
+      }
+    }
+  }
+`);
+
+let forAwaitVisitor = {
+  noScope: true,
+
+  Identifier(path, replacements) {
+    if (path.node.name in replacements) {
+      path.replaceInline(replacements[path.node.name]);
+    }
+  },
+
+  CallExpression(path, replacements) {
+    let callee = path.node.callee;
+
+    // if no await wrapping is being applied, unwrap the call expression
+    if (t.isIdentifier(callee) && callee.name === "AWAIT" && !replacements.AWAIT) {
+      path.replaceWith(path.node.arguments[0]);
+    }
+  }
+};
+
+export default function (path, helpers) {
+  let { node, scope, parent } = path;
+
+  let stepKey = scope.generateUidIdentifier("step");
+  let stepValue = scope.generateUidIdentifier("value");
+  let left = node.left;
+  let declar;
+
+  if (t.isIdentifier(left) || t.isPattern(left) || t.isMemberExpression(left)) {
+    // for await (i of test), for await ({ i } of test)
+    declar = t.expressionStatement(t.assignmentExpression("=", left, stepValue));
+  } else if (t.isVariableDeclaration(left)) {
+    // for await (let i of test)
+    declar = t.variableDeclaration(left.kind, [
+      t.variableDeclarator(left.declarations[0].id, stepValue)
+    ]);
+  }
+
+  let template = buildForAwait();
+
+  traverse(template, forAwaitVisitor, null, {
+    ITERATOR_HAD_ERROR_KEY: scope.generateUidIdentifier("didIteratorError"),
+    ITERATOR_COMPLETION: scope.generateUidIdentifier("iteratorNormalCompletion"),
+    ITERATOR_ERROR_KEY: scope.generateUidIdentifier("iteratorError"),
+    ITERATOR_KEY: scope.generateUidIdentifier("iterator"),
+    GET_ITERATOR: helpers.getAsyncIterator,
+    OBJECT: node.right,
+    STEP_VALUE: stepValue,
+    STEP_KEY: stepKey,
+    AWAIT: helpers.wrapAwait
+  });
+
+  // remove generator function wrapper
+  template = template.body.body;
+
+  let isLabeledParent = t.isLabeledStatement(parent);
+  let tryBody = template[3].block.body;
+  let loop = tryBody[0];
+
+  if (isLabeledParent) {
+    tryBody[0] = t.labeledStatement(parent.label, loop);
+  }
+
+  return {
+    replaceParent: isLabeledParent,
+    node: template,
+    declar,
+    loop
+  };
+}

packages/babel-helper-remap-async-to-generator/src/index.js

@@ -4,6 +4,7 @@ import type { NodePath } from "babel-traverse";
 import nameFunction from "babel-helper-function-name";
 import template from "babel-template";
 import * as t from "babel-types";
+import rewriteForAwait from "./for-await";
 
 let buildWrapper = template(`
   (() => {
@@ -25,15 +26,54 @@ let namedBuildWrapper = template(`
 `);
 
 let awaitVisitor = {
-  ArrowFunctionExpression(path) {
-    if (!path.node.async) {
+  Function(path) {
+    if (path.isArrowFunctionExpression() && !path.node.async) {
       path.arrowFunctionToShadowed();
+      return;
     }
+    path.skip();
   },
 
-  AwaitExpression({ node }) {
+  AwaitExpression({ node }, { wrapAwait }) {
     node.type = "YieldExpression";
+    if (wrapAwait) {
+      node.argument = t.callExpression(wrapAwait, [node.argument]);
+    }
+  },
+
+  ForAwaitStatement(path, { file, wrapAwait }) {
+    let { node } = path;
+
+    let build = rewriteForAwait(path, {
+      getAsyncIterator: file.addHelper("asyncIterator"),
+      wrapAwait
+    });
+
+    let { declar, loop } = build;
+    let block = loop.body;
+
+    // ensure that it's a block so we can take all its statements
+    path.ensureBlock();
+
+    // add the value declaration to the new loop body
+    if (declar) {
+      block.body.push(declar);
+    }
+
+    // push the rest of the original loop body onto our new body
+    block.body = block.body.concat(node.body.body);
+
+    t.inherits(loop, node);
+    t.inherits(loop.body, node.body);
+
+    if (build.replaceParent) {
+      path.parentPath.replaceWithMultiple(build.node);
+      path.remove();
+    } else {
+      path.replaceWithMultiple(build.node);
+    }
   }
+
 };
 
 function classOrObjectMethod(path: NodePath, callId: Object) {
@@ -110,15 +150,15 @@ function plainFunction(path: NodePath, callId: Object) {
   }
 }
 
-export default function (path: NodePath, callId: Object) {
-  let node = path.node;
-  if (node.generator) return;
-
-  path.traverse(awaitVisitor);
+export default function (path: NodePath, file: Object, helpers: Object) {
+  path.get("body").traverse(awaitVisitor, {
+    file,
+    wrapAwait: helpers.wrapAwait
+  });
 
   if (path.isClassMethod() || path.isObjectMethod()) {
-    return classOrObjectMethod(path, callId);
+    classOrObjectMethod(path, helpers.wrapAsync);
   } else {
-    return plainFunction(path, callId);
+    plainFunction(path, helpers.wrapAsync);
   }
 }

packages/babel-helpers/src/helpers.js

@@ -57,6 +57,158 @@ helpers.jsx = template(`
   })()
 `);
 
+helpers.asyncIterator = template(`
+  (function (iterable) {
+    if (typeof Symbol === "function") {
+      if (Symbol.asyncIterator) {
+        var method = iterable[Symbol.asyncIterator];
+        if (method != null) return method.call(iterable);
+      }
+      if (Symbol.iterator) {
+        return iterable[Symbol.iterator]();
+      }
+    }
+    throw new TypeError("Object is not async iterable");
+  })
+`);
+
+helpers.asyncGenerator = template(`
+  (function () {
+    function AwaitValue(value) {
+      this.value = value;
+    }
+
+    function AsyncGenerator(gen) {
+      var front, back;
+
+      function send(key, arg) {
+        return new Promise(function (resolve, reject) {
+          var request = {
+            key: key,
+            arg: arg,
+            resolve: resolve,
+            reject: reject,
+            next: null
+          };
+
+          if (back) {
+            back = back.next = request;
+          } else {
+            front = back = request;
+            resume(key, arg);
+          }
+        });
+      }
+
+      function resume(key, arg) {
+        try {
+          var result = gen[key](arg)
+          var value = result.value;
+          if (value instanceof AwaitValue) {
+            Promise.resolve(value.value).then(
+              function (arg) { resume("next", arg); },
+              function (arg) { resume("throw", arg); });
+          } else {
+            settle(result.done ? "return" : "normal", result.value);
+          }
+        } catch (err) {
+          settle("throw", err);
+        }
+      }
+
+      function settle(type, value) {
+        switch (type) {
+          case "return":
+            front.resolve({ value: value, done: true });
+            break;
+          case "throw":
+            front.reject(value);
+            break;
+          default:
+            front.resolve({ value: value, done: false });
+            break;
+        }
+
+        front = front.next;
+        if (front) {
+          resume(front.key, front.arg);
+        } else {
+          back = null;
+        }
+      }
+
+      this._invoke = send;
+
+      // Hide "return" method if generator return is not supported
+      if (typeof gen.return !== "function") {
+        this.return = undefined;
+      }
+    }
+
+    if (typeof Symbol === "function" && Symbol.asyncIterator) {
+      AsyncGenerator.prototype[Symbol.asyncIterator] = function () { return this; };
+    }
+
+    AsyncGenerator.prototype.next = function (arg) { return this._invoke("next", arg); };
+    AsyncGenerator.prototype.throw = function (arg) { return this._invoke("throw", arg); };
+    AsyncGenerator.prototype.return = function (arg) { return this._invoke("return", arg); };
+
+    return {
+      wrap: function (fn) {
+        return function () {
+          return new AsyncGenerator(fn.apply(this, arguments));
+        };
+      },
+      await: function (value) {
+        return new AwaitValue(value);
+      }
+    };
+
+  })()
+`);
+
+helpers.asyncGeneratorDelegate = template(`
+  (function (inner, awaitWrap) {
+    var iter = {}, waiting = false;
+
+    function pump(key, value) {
+      waiting = true;
+      value = new Promise(function (resolve) { resolve(inner[key](value)); });
+      return { done: false, value: awaitWrap(value) };
+    };
+
+    if (typeof Symbol === "function" && Symbol.iterator) {
+      iter[Symbol.iterator] = function () { return this; };
+    }
+
+    iter.next = function (value) {
+      if (waiting) {
+        waiting = false;
+        return value;
+      }
+      return pump("next", value);
+    };
+
+    if (typeof inner.throw === "function") {
+      iter.throw = function (value) {
+        if (waiting) {
+          waiting = false;
+          throw value;
+        }
+        return pump("throw", value);
+      };
+    }
+
+    if (typeof inner.return === "function") {
+      iter.return = function (value) {
+        return pump("return", value);
+      };
+    }
+
+    return iter;
+  })
+`);
+
 helpers.asyncToGenerator = template(`
   (function (fn) {
     return function () {
@@ -88,7 +240,6 @@ helpers.asyncToGenerator = template(`
   })
 `);
 
-
 helpers.classCallCheck = template(`
   (function (instance, Constructor) {
     if (!(instance instanceof Constructor)) {

packages/babel-plugin-syntax-async-generators/src/index.js

@@ -1,7 +1,7 @@
 export default function () {
   return {
     manipulateOptions(opts, parserOpts) {
-      parserOpts.plugins.push("asyncGenerators");
+      parserOpts.plugins.push("asyncFunctions", "asyncGenerators");
     }
   };
 }

packages/babel-plugin-transform-async-generator-functions/.npmignore

@@ -0,0 +1,4 @@
+node_modules
+*.log
+src
+test