From 06a58f228cd587df4761daf93ec9a661da84a0ce Mon Sep 17 00:00:00 2001
From: Sebastian McKenzie <sebmck@gmail.com>
Date: Sat, 28 Mar 2015 03:59:51 +1100
Subject: [PATCH] add TraversalPath#replaceInline, fix traversal path not
 updating sibling keys correctly

---
 src/babel/traversal/context.js    | 44 +++++++++++++++++--------
 src/babel/traversal/path/index.js | 53 +++++++++++++++++++++++++------
 src/babel/util.js                 |  2 +-
 3 files changed, 76 insertions(+), 23 deletions(-)

diff --git a/src/babel/traversal/context.js b/src/babel/traversal/context.js
index 928c651279..1642bc8b6f 100644
--- a/src/babel/traversal/context.js
+++ b/src/babel/traversal/context.js
@@ -14,29 +14,47 @@ export default class TraversalContext {
     return TraversalPath.get(this.parentPath, this, node, obj, key);
   }
 
-  visit(node, key) {
-    var nodes = node[key];
-    if (!nodes) return;
-
-    if (!Array.isArray(nodes)) {
-      return this.create(node, node, key).visit();
-    }
-
+  visitMultiple(nodes, node, key) {
     // nothing to traverse!
-    if (nodes.length === 0) {
-      return;
-    }
+    if (nodes.length === 0) return false;
 
-    var queue = [];
+    var queue = this.queue = [];
+    var stop  = false;
+    var done  = [];
 
+    // build up initial queue
     for (let i = 0; i < nodes.length; i++) {
       if (nodes[i]) queue.push(this.create(node, nodes, i));
     }
 
+    // visit the queue
     for (let i = 0; i < queue.length; i++) {
       if (queue[i].visit()) {
-        return true;
+        stop = true;
+        break;
       }
     }
+
+    // clear context from queued paths
+    for (let i = 0; i < queue.length; i++) {
+      //queue[i].clearContext();
+    }
+
+    return stop;
+  }
+
+  visitSingle(node, key) {
+    return this.create(node, node, key).visit();
+  }
+
+  visit(node, key) {
+    var nodes = node[key];
+    if (!nodes) return;
+
+    if (Array.isArray(nodes)) {
+      return this.visitMultiple(nodes, node, key);
+    } else {
+      return this.visitSingle(node, key);
+    }
   }
 }
diff --git a/src/babel/traversal/path/index.js b/src/babel/traversal/path/index.js
index aaa7227058..d6a164c01c 100644
--- a/src/babel/traversal/path/index.js
+++ b/src/babel/traversal/path/index.js
@@ -78,6 +78,12 @@ export default class TraversalPath {
     return ourScope;
   }
 
+  queueNode(path) {
+    if (this.context) {
+      this.context.queue.push(path);
+    }
+  }
+
   insertBefore(nodes) {
     this.checkNodes(nodes);
 
@@ -89,15 +95,24 @@ export default class TraversalPath {
     }
   }
 
+  _containerInsertAfter(nodes) {
+    this.updateSiblingKeys(this.key + 1, nodes.length);
+    for (var i = 0; i < nodes.length; i++) {
+      var to = this.key + 1 + i;
+      this.container.splice(to, 0, nodes[i]);
+
+      if (this.context) {
+        this.queueNode(this.context.create(this.parent, this.container, to));
+      }
+    }
+  }
+
   insertAfter(nodes) {
     this.checkNodes(nodes);
 
     if (this.isPreviousType("Statement")) {
       if (Array.isArray(this.container)) {
-        for (var i = 0; i < nodes.length; i++) {
-          this.container.splice(this.key + 1 + i, 0, nodes[i]);
-        }
-        this.updateSiblingKeys(this.key + nodes.length, nodes.length);
+        this._containerInsertAfter(nodes);
       } else if (includes(t.STATEMENT_OR_BLOCK_KEYS, this.key) && !t.isBlockStatement(this.container)) {
         this.container[this.key] = t.blockStatement(nodes);
       } else {
@@ -139,6 +154,10 @@ export default class TraversalPath {
     this.scope = TraversalPath.getScope(this, this.context && this.context.scope, file);
   }
 
+  clearContext() {
+    this.context = null;
+  }
+
   setContext(parentPath, context, key, file?) {
     this.shouldSkip = false;
     this.shouldStop = false;
@@ -200,18 +219,34 @@ export default class TraversalPath {
     throw new Error("Don't use `path.node = newNode;`, use `path.replaceWith(newNode)` or `path.replaceWithMultiple([newNode])`");
   }
 
-  replaceWithMultiple(nodes: Array<Object>) {
+  replaceInline(nodes) {
+    if (Array.isArray(nodes)) {
+      if (Array.isArray(this.container)) {
+        this._verifyNodeList("replaceInline", nodes);
+        this._containerInsertAfter(nodes);
+        return this.remove();
+      } else {
+        return this.replaceWithMultiple(nodes);
+      }
+    } else {
+      return this.replaceWith(nodes);
+    }
+  }
+
+  _verifyNodeList(key, nodes) {
     if (nodes.indexOf(this.node) >= 0) {
-      // todo: check for inclusion of current node in `nodes` and yell at the user if it's in there and tell them to use `insertBefore` or `insertAfter`
+      // todo: possibly check for inclusion of current node and yell at the user as it's an anti-pattern
     }
 
     for (var i = 0; i < nodes.length; i++) {
       var node = nodes[i];
-      if (!node) throw new Error(`Falsy node passed to \`path.replaceWithMultiple()\` with the index of ${i}`);
+      if (!node) throw new Error(`Falsy node passed to \`path.${key}()\` with the index of ${i}`);
     }
+  }
 
-    t.inherits(nodes[0], this.node);
-
+  replaceWithMultiple(nodes: Array<Object>) {
+    this._verifyNodeList("replaceWithMultiple", nodes);
+    t.inheritsComments(nodes[0], this.node);
     this.container[this.key] = null;
     this.insertAfter(nodes);
     if (!this.node) this.remove();
diff --git a/src/babel/util.js b/src/babel/util.js
index 61bccf1410..968b08d763 100644
--- a/src/babel/util.js
+++ b/src/babel/util.js
@@ -76,7 +76,7 @@ var templateVisitor = {
     }
     if (t.isIdentifier(node) && has(nodes, node.name)) {
       this.skip();
-      return nodes[node.name];
+      this.replaceInline(nodes[node.name]);
     }
   }
 };