From 0717eaddce47c6fd8f05312c346a11c5cf61f60f Mon Sep 17 00:00:00 2001
From: Sebastian McKenzie <sebmck@gmail.com>
Date: Sun, 3 May 2015 23:35:39 +0100
Subject: [PATCH] normalise Program replacement nodes

---
 .../transformers/other/regenerator.js           |  2 +-
 src/babel/traversal/path/index.js               | 17 +++++++++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/babel/transformation/transformers/other/regenerator.js b/src/babel/transformation/transformers/other/regenerator.js
index d330eeacdb..7c109a92a5 100644
--- a/src/babel/transformation/transformers/other/regenerator.js
+++ b/src/babel/transformation/transformers/other/regenerator.js
@@ -9,6 +9,6 @@ export var Program = {
   enter(ast) {
     regenerator.transform(ast);
     this.stop();
-    return ast; // force a checkPath, this really needs to be optimised
+    this.checkSelf();
   }
 };
diff --git a/src/babel/traversal/path/index.js b/src/babel/traversal/path/index.js
index 0be1359729..166e15e964 100644
--- a/src/babel/traversal/path/index.js
+++ b/src/babel/traversal/path/index.js
@@ -118,7 +118,7 @@ export default class TraversalPath {
       } else if (this.isStatementOrBlock()) {
         if (this.node) nodes.push(this.node);
         this.container[this.key] = t.blockStatement(nodes);
-        this.checkPaths(this);
+        this.checkSelf();
       } else {
         throw new Error("We don't know what to do with this node type. We were previously a Statement but we can't fit in here?");
       }
@@ -216,7 +216,7 @@ export default class TraversalPath {
       } else if (this.isStatementOrBlock()) {
         if (this.node) nodes.unshift(this.node);
         this.container[this.key] = t.blockStatement(nodes);
-        this.checkPaths(this);
+        this.checkSelf();
       } else {
         throw new Error("We don't know what to do with this node type. We were previously a Statement but we can't fit in here?");
       }
@@ -509,6 +509,11 @@ export default class TraversalPath {
       throw new Error("You passed `path.replaceWith()` a falsy node, use `path.remove()` instead");
     }
 
+    // normalise inserting an entire AST
+    if (t.isProgram(replacement)) {
+      replacement = replacement.body;
+    }
+
     if (Array.isArray(replacement)) {
       if (whateverAllowed) {
         return this.replaceWithMultiple(replacement);
@@ -545,6 +550,14 @@ export default class TraversalPath {
     // potentially create new scope
     this.setScope();
 
+    this.checkSelf();
+  }
+
+  /**
+   * Description
+   */
+
+  checkSelf() {
     this.checkPaths(this);
   }