add iterator.return to for-of breaks - fixes #838

lib/babel/transformation/transformers/es6/for-of.js

@@ -37,6 +37,29 @@ exports.ForOfStatement = function (node, parent, scope, file) {
   return loop;
 };
 
+var breakVisitor = {
+  enter: function (node, parent, scope, state) {
+    if (t.isLoop(node)) {
+      state.ignoreLabeless = true;
+      scope.traverse(node, breakVisitor, state);
+      state.ignoreLabeless = false;
+      return this.skip();
+    }
+
+    if (t.isBreakStatement(node)) {
+      if (!node.label && state.ignoreLabeless) return;
+      if (node.label && node.label.name !== state.label) return;
+
+      var ret = t.expressionStatement(
+        t.callExpression(t.memberExpression(state.iteratorKey, t.identifier("return")), [])
+      );
+      if (state.wrapReturn) ret = state.wrapReturn(ret);
+
+      return [ret, node];
+    }
+  }
+};
+
 var loose = function (node, parent, scope, file) {
   var left = node.left;
   var declar, id;
@@ -54,9 +77,12 @@ var loose = function (node, parent, scope, file) {
     throw file.errorWithNode(left, messages.get("unknownForHead", left.type));
   }
 
+  var iteratorKey = scope.generateUidIdentifier("iterator");
+  var isArrayKey  = scope.generateUidIdentifier("isArray");
+
   var loop = util.template("for-of-loose", {
-    LOOP_OBJECT:  scope.generateUidIdentifier("iterator"),
-    IS_ARRAY:     scope.generateUidIdentifier("isArray"),
+    LOOP_OBJECT:  iteratorKey,
+    IS_ARRAY:     isArrayKey,
     OBJECT:       node.right,
     INDEX:        scope.generateUidIdentifier("i"),
     ID:           id
@@ -68,6 +94,18 @@ var loose = function (node, parent, scope, file) {
     loop.body.body.shift();
   }
 
+  //
+
+  scope.traverse(node, breakVisitor, {
+    iteratorKey: iteratorKey,
+    wrapReturn:  function (node) {
+      return t.ifStatement(t.unaryExpression("!", isArrayKey, true), node);
+    },
+    label:       t.isLabeledStatement(parent) && parent.label.name
+  });
+
+  //
+
   return {
     declar: declar,
     loop:   loop
@@ -93,12 +131,25 @@ var spec = function (node, parent, scope, file) {
     throw file.errorWithNode(left, messages.get("unknownForHead", left.type));
   }
 
+  //
+
+  var iteratorKey = scope.generateUidIdentifier("iterator");
+
   var loop = util.template("for-of", {
-    ITERATOR_KEY: scope.generateUidIdentifier("iterator"),
+    ITERATOR_KEY: iteratorKey,
     STEP_KEY:     stepKey,
     OBJECT:       node.right
   });
 
+  //
+
+  scope.traverse(node, breakVisitor, {
+    iteratorKey: iteratorKey,
+    label:       t.isLabeledStatement(parent) && parent.label.name
+  });
+
+  //
+
   return {
     declar: declar,
     loop:   loop

test/fixtures/transformation/es6-for-of-loose/break/actual.js

@@ -0,0 +1,29 @@
+// labels
+
+foo: for (let x of foo()) {
+  while (true) {
+    break foo;
+  }
+}
+
+foo: for (let x of foo()) {
+  while (true) {
+    break;
+  }
+}
+
+foo: for (let x of foo()) {
+  break foo;
+}
+
+// basic
+
+for (let x of foo()) {
+  break;
+}
+
+for (let x of foo()) {
+  while (true) {
+    break;
+  }
+}

test/fixtures/transformation/es6-for-of-loose/break/expected.js

@@ -0,0 +1,100 @@
+"use strict";
+
+// labels
+
+foo: for (var _iterator = foo(), _isArray = Array.isArray(_iterator), _i = 0, _iterator = _isArray ? _iterator : _iterator[Symbol.iterator]();;) {
+  var _ref;
+
+  if (_isArray) {
+    if (_i >= _iterator.length) break;
+    _ref = _iterator[_i++];
+  } else {
+    _i = _iterator.next();
+    if (_i.done) break;
+    _ref = _i.value;
+  }
+
+  var x = _ref;
+
+  while (true) {
+    if (!_isArray) _iterator["return"]();
+
+    break foo;
+  }
+}
+
+foo: for (var _iterator2 = foo(), _isArray2 = Array.isArray(_iterator2), _i2 = 0, _iterator2 = _isArray2 ? _iterator2 : _iterator2[Symbol.iterator]();;) {
+  var _ref2;
+
+  if (_isArray2) {
+    if (_i2 >= _iterator2.length) break;
+    _ref2 = _iterator2[_i2++];
+  } else {
+    _i2 = _iterator2.next();
+    if (_i2.done) break;
+    _ref2 = _i2.value;
+  }
+
+  var x = _ref2;
+
+  while (true) {
+    break;
+  }
+}
+
+foo: for (var _iterator3 = foo(), _isArray3 = Array.isArray(_iterator3), _i3 = 0, _iterator3 = _isArray3 ? _iterator3 : _iterator3[Symbol.iterator]();;) {
+  var _ref3;
+
+  if (_isArray3) {
+    if (_i3 >= _iterator3.length) break;
+    _ref3 = _iterator3[_i3++];
+  } else {
+    _i3 = _iterator3.next();
+    if (_i3.done) break;
+    _ref3 = _i3.value;
+  }
+
+  var x = _ref3;
+  if (!_isArray3) _iterator3["return"]();
+
+  break foo;
+}
+
+// basic
+
+for (var _iterator4 = foo(), _isArray4 = Array.isArray(_iterator4), _i4 = 0, _iterator4 = _isArray4 ? _iterator4 : _iterator4[Symbol.iterator]();;) {
+  var _ref4;
+
+  if (_isArray4) {
+    if (_i4 >= _iterator4.length) break;
+    _ref4 = _iterator4[_i4++];
+  } else {
+    _i4 = _iterator4.next();
+    if (_i4.done) break;
+    _ref4 = _i4.value;
+  }
+
+  var x = _ref4;
+  if (!_isArray4) _iterator4["return"]();
+
+  break;
+}
+
+for (var _iterator5 = foo(), _isArray5 = Array.isArray(_iterator5), _i5 = 0, _iterator5 = _isArray5 ? _iterator5 : _iterator5[Symbol.iterator]();;) {
+  var _ref5;
+
+  if (_isArray5) {
+    if (_i5 >= _iterator5.length) break;
+    _ref5 = _iterator5[_i5++];
+  } else {
+    _i5 = _iterator5.next();
+    if (_i5.done) break;
+    _ref5 = _i5.value;
+  }
+
+  var x = _ref5;
+
+  while (true) {
+    break;
+  }
+}

test/fixtures/transformation/es6-for-of/break/actual.js

@@ -0,0 +1,29 @@
+// labels
+
+foo: for (let x of foo()) {
+  while (true) {
+    break foo;
+  }
+}
+
+foo: for (let x of foo()) {
+  while (true) {
+    break;
+  }
+}
+
+foo: for (let x of foo()) {
+  break foo;
+}
+
+// basic
+
+for (let x of foo()) {
+  break;
+}
+
+for (let x of foo()) {
+  while (true) {
+    break;
+  }
+}

test/fixtures/transformation/es6-for-of/break/expected.js

@@ -0,0 +1,47 @@
+"use strict";
+
+// labels
+
+foo: for (var _iterator = foo()[Symbol.iterator](), _step; !(_step = _iterator.next()).done;) {
+  var x = _step.value;
+
+  while (true) {
+    _iterator["return"]();
+
+    break foo;
+  }
+}
+
+foo: for (var _iterator2 = foo()[Symbol.iterator](), _step2; !(_step2 = _iterator2.next()).done;) {
+  var x = _step2.value;
+
+  while (true) {
+    break;
+  }
+}
+
+foo: for (var _iterator3 = foo()[Symbol.iterator](), _step3; !(_step3 = _iterator3.next()).done;) {
+  var x = _step3.value;
+
+  _iterator3["return"]();
+
+  break foo;
+}
+
+// basic
+
+for (var _iterator4 = foo()[Symbol.iterator](), _step4; !(_step4 = _iterator4.next()).done;) {
+  var x = _step4.value;
+
+  _iterator4["return"]();
+
+  break;
+}
+
+for (var _iterator5 = foo()[Symbol.iterator](), _step5; !(_step5 = _iterator5.next()).done;) {
+  var x = _step5.value;
+
+  while (true) {
+    break;
+  }
+}