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