From 9b12f799f7d489afe5dfd179aa4f93a714c2cecd Mon Sep 17 00:00:00 2001
From: Sebastian McKenzie <sebmck@gmail.com>
Date: Wed, 20 May 2015 10:07:29 +0100
Subject: [PATCH] clean up system module formatter hoisting visitor and allow
 contextual replacement of variable declarations with expressions in for head
 positions - fixes #1570

---
 src/babel/transformation/modules/system.js | 60 +++++++++++-----------
 src/babel/traversal/path/index.js          | 16 +++++-
 2 files changed, 44 insertions(+), 32 deletions(-)

diff --git a/src/babel/transformation/modules/system.js b/src/babel/transformation/modules/system.js
index c8ff4c4275..542b84f228 100644
--- a/src/babel/transformation/modules/system.js
+++ b/src/babel/transformation/modules/system.js
@@ -8,41 +8,39 @@ import map from "lodash/collection/map";
 import * as t from "../../types";
 
 var hoistVariablesVisitor = {
-  enter(node, parent, scope, state) {
-    if (t.isFunction(node)) {
-      // nothing inside is accessible
-      return this.skip();
+  Function() {
+    // nothing inside is accessible
+    this.skip();
+  },
+
+  VariableDeclaration(node, parent, scope, state) {
+    if (node.kind !== "var" && !t.isProgram(parent)) { // let, const
+      // can't be accessed
+      return;
     }
 
-    if (t.isVariableDeclaration(node)) {
-      if (node.kind !== "var" && !t.isProgram(parent)) { // let, const
-        // can't be accessed
-        return;
+    // ignore block hoisted nodes as these can be left in
+    if (state.formatter._canHoist(node)) return;
+
+    var nodes = [];
+
+    for (var i = 0; i < node.declarations.length; i++) {
+      var declar = node.declarations[i];
+      state.hoistDeclarators.push(t.variableDeclarator(declar.id));
+      if (declar.init) {
+        // no initializer so we can just hoist it as-is
+        var assign = t.expressionStatement(t.assignmentExpression("=", declar.id, declar.init));
+        nodes.push(assign);
       }
-
-      // ignore block hoisted nodes as these can be left in
-      if (state.formatter._canHoist(node)) return;
-
-      var nodes = [];
-
-      for (var i = 0; i < node.declarations.length; i++) {
-        var declar = node.declarations[i];
-        state.hoistDeclarators.push(t.variableDeclarator(declar.id));
-        if (declar.init) {
-          // no initializer so we can just hoist it as-is
-          var assign = t.expressionStatement(t.assignmentExpression("=", declar.id, declar.init));
-          nodes.push(assign);
-        }
-      }
-
-      // for (var i in test)
-      // for (var i = 0;;)
-      if (t.isFor(parent) && parent.left === node) {
-        return node.declarations[0].id;
-      }
-
-      return nodes;
     }
+
+    // for (var i in test)
+    // for (var i = 0;;)
+    if (t.isFor(parent) && (parent.left === node || parent.init === node)) {
+      return node.declarations[0].id;
+    }
+
+    return nodes;
   }
 };
 
diff --git a/src/babel/traversal/path/index.js b/src/babel/traversal/path/index.js
index 2c52c71f67..5f6507ca18 100644
--- a/src/babel/traversal/path/index.js
+++ b/src/babel/traversal/path/index.js
@@ -708,7 +708,7 @@ export default class TraversalPath {
     }
 
     // replacing a statement with an expression so wrap it in an expression statement
-    if (this.isPreviousType("Statement") && t.isExpression(replacement)) {
+    if (this.isPreviousType("Statement") && t.isExpression(replacement) && !this.canHaveVariableDeclarationOrExpression()) {
       replacement = t.expressionStatement(replacement);
     }
 
@@ -728,6 +728,20 @@ export default class TraversalPath {
     this.setScope();
   }
 
+  /**
+   * This checks whether or now we're in one of the following positions:
+   *
+   *   for (KEY in right);
+   *   for (KEY;;);
+   *
+   * This is because these spots allow VariableDeclarations AND normal expressions so we need to tell the
+   * path replacement that it's ok to replace this with an expression.
+   */
+
+   canHaveVariableDeclarationOrExpression() {
+      return (this.key === "init" || this.key === "left") && this.parentPath.isFor();
+   }
+
   /**
    * Description
    */