From cc611cb71ca225e24e65158be0cc0e18357957d4 Mon Sep 17 00:00:00 2001
From: Ben Newman <bn@cs.stanford.edu>
Date: Thu, 14 May 2015 13:34:43 -0500
Subject: [PATCH] Invoke Regenerator on async/generator Function nodes, not
 whole AST.

This should help (dramatically?) with
https://github.com/babel/babel/issues/1486#issuecomment-101491605,
although I'm not sure how to run the benchmarks myself.
---
 .../transformers/other/regenerator.js         | 64 ++++++++++++++++++-
 1 file changed, 61 insertions(+), 3 deletions(-)

diff --git a/src/babel/transformation/transformers/other/regenerator.js b/src/babel/transformation/transformers/other/regenerator.js
index e7656aa6f0..d2a0112459 100644
--- a/src/babel/transformation/transformers/other/regenerator.js
+++ b/src/babel/transformation/transformers/other/regenerator.js
@@ -1,11 +1,69 @@
 import regenerator from "regenerator";
 import * as t from "../../../types";
+import { NodePath } from "ast-types";
 
 export var metadata = {
   group: "regenerator"
 };
 
-export function Program(ast) {
-  regenerator.transform(ast);
-  this.stop();
+export function Func/*tion*/(node) {
+  if (node.async || node.generator) {
+    // Although this code transforms only the subtree rooted at the given
+    // Function node, that node might contain other generator functions
+    // that will also be transformed. It might help performance to ignore
+    // nested functions, and rely on the traversal to visit them later,
+    // but that's a small optimization. Starting here instead of at the
+    // root of the AST is the key optimization, since huge async/generator
+    // functions are relatively rare.
+    regenerator.transform(convertTraversalPathToNodePath(this));
+  }
+}
+
+// Given a TraversalPath, return a NodePath that includes full ancestry
+// information (up to and including the Program node). This is complicated
+// by having to include intermediate objects like blockStatement.body
+// arrays, in addition to Node objects.
+function convertTraversalPathToNodePath(path) {
+  var programNode;
+  var keysAlongPath = [];
+
+  while (path) {
+    var pp = path.parentPath;
+    var parentNode = pp && pp.node;
+    if (parentNode) {
+      keysAlongPath.push(path.key);
+
+      if (parentNode !== path.container) {
+        var found = Object.keys(parentNode).some(containerKey => {
+          if (parentNode[containerKey] === path.container) {
+            keysAlongPath.push(containerKey);
+            return true;
+          }
+        });
+
+        if (!found) {
+          throw new Error("Failed to find container object in parent node");
+        }
+      }
+
+      if (t.isProgram(parentNode)) {
+        programNode = parentNode;
+        break;
+      }
+    }
+
+    path = pp;
+  }
+
+  if (!programNode) {
+    throw new Error("Failed to find root Program node");
+  }
+
+  var nodePath = new NodePath(programNode);
+
+  while (keysAlongPath.length > 0) {
+    nodePath = nodePath.get(keysAlongPath.pop());
+  }
+
+  return nodePath;
 }