for-of optimization on arrays/known functions that return arrays (#4747)

packages/babel-core/test/fixtures/transformation/misc/regression-1169/expected.js

@@ -2,31 +2,11 @@ function foo() {
   var input = ['a', 'b', 'c'];
   var output = {};
 
-  var _iteratorNormalCompletion = true;
-  var _didIteratorError = false;
-  var _iteratorError = undefined;
-
-  try {
-    for (var _iterator = input[Symbol.iterator](), _step; !(_iteratorNormalCompletion = (_step = _iterator.next()).done); _iteratorNormalCompletion = true) {
-      var c = _step.value;
-
-      var name = c;
-      output[name] = name;
-    }
-  } catch (err) {
-    _didIteratorError = true;
-    _iteratorError = err;
-  } finally {
-    try {
-      if (!_iteratorNormalCompletion && _iterator.return) {
-        _iterator.return();
-      }
-    } finally {
-      if (_didIteratorError) {
-        throw _iteratorError;
-      }
-    }
+  for (var _i = 0; _i < input.length; _i++) {
+    var c = input[_i];
+    var name = c;
+    output[name] = name;
   }
 
   return output;
-}
+}

packages/babel-plugin-transform-es2015-for-of/src/index.js

@@ -1,4 +1,9 @@
 export default function ({ messages, template, types: t }) {
+  const isArrayFrom = t.buildMatchMemberExpression("Array.from");
+  const isObjectKeys = t.buildMatchMemberExpression("Object.keys");
+  const isObjectValues = t.buildMatchMemberExpression("Object.values");
+  const isObjectEntries = t.buildMatchMemberExpression("Object.entries");
+
   const buildForOfArray = template(`
     for (var KEY = 0; KEY < ARR.length; KEY++) BODY;
   `);
@@ -89,16 +94,48 @@ export default function ({ messages, template, types: t }) {
     return nodes;
   }
 
+  function replaceWithArray(path) {
+    if (path.parentPath.isLabeledStatement()) {
+      path.parentPath.replaceWithMultiple(_ForOfStatementArray(path));
+      return true;
+    } else {
+      path.replaceWithMultiple(_ForOfStatementArray(path));
+      return true;
+    }
+    return false;
+  }
+
+  function optimize(path, right) {
+    if (right.isArrayExpression() || right.isGenericType("Array")) {
+      return replaceWithArray(path);
+    } else if (right.isIdentifier() && right.isPure()) {
+      const binding = path.scope.getBinding(right.node.name);
+      return optimize(path, binding.path.get("init"));
+    } else if (right.isCallExpression() && (
+      isArrayFrom(right.get("callee").node) ||
+      isObjectKeys(right.get("callee").node) ||
+      isObjectValues(right.get("callee").node) ||
+      isObjectEntries(right.get("callee").node)
+      )
+    ) {
+      const initPath = right === path.get("right") ? path : right.find((p) => p.isStatement());
+      const uid = path.scope.generateUidIdentifierBasedOnNode(right.node);
+      initPath.insertBefore(
+        t.variableDeclaration("const", [
+          t.variableDeclarator(uid, right.node),
+        ])
+      );
+      right.replaceWith(uid);
+      return replaceWithArray(path);
+    }
+    return false;
+  }
 
   return {
     visitor: {
       ForOfStatement(path, state) {
-        if (path.get("right").isArrayExpression()) {
-          if (path.parentPath.isLabeledStatement()) {
-            return path.parentPath.replaceWithMultiple(_ForOfStatementArray(path));
-          } else {
-            return path.replaceWithMultiple(_ForOfStatementArray(path));
-          }
+        if (optimize(path, path.get("right"))) {
+          return;
         }
 
         let callback = spec;

packages/babel-plugin-transform-es2015-for-of/test/fixtures/opt/array-binding/actual.js

@@ -0,0 +1,5 @@
+const x = [];
+for (const y of x) {}
+const arr = Object.entries(x);
+for (const y of arr) {}
+

packages/babel-plugin-transform-es2015-for-of/test/fixtures/opt/array-binding/expected.js

@@ -0,0 +1,11 @@
+const x = [];
+for (var _i = 0; _i < x.length; _i++) {
+  const y = x[_i];
+}
+
+const _Object$entries = Object.entries(x);
+
+const arr = _Object$entries;
+for (var _i2 = 0; _i2 < arr.length; _i2++) {
+  const y = arr[_i2];
+}

packages/babel-plugin-transform-es2015-for-of/test/fixtures/opt/built-ins/actual.js

@@ -0,0 +1,5 @@
+for (const y of []) {}
+for (const y of Array.from(x)) {}
+for (const y of Object.keys(x)) {}
+for (const y of Object.values(x)) {}
+for (const y of Object.entries(x)) {}

packages/babel-plugin-transform-es2015-for-of/test/fixtures/opt/built-ins/expected.js

@@ -0,0 +1,28 @@
+var _arr = [];
+for (var _i = 0; _i < _arr.length; _i++) {
+  const y = _arr[_i];
+}
+
+const _Array$from = Array.from(x);
+
+for (var _i2 = 0; _i2 < _Array$from.length; _i2++) {
+  const y = _Array$from[_i2];
+}
+
+const _Object$keys = Object.keys(x);
+
+for (var _i3 = 0; _i3 < _Object$keys.length; _i3++) {
+  const y = _Object$keys[_i3];
+}
+
+const _Object$values = Object.values(x);
+
+for (var _i4 = 0; _i4 < _Object$values.length; _i4++) {
+  const y = _Object$values[_i4];
+}
+
+const _Object$entries = Object.entries(x);
+
+for (var _i5 = 0; _i5 < _Object$entries.length; _i5++) {
+  const y = _Object$entries[_i5];
+}

packages/babel-plugin-transform-es2015-for-of/test/fixtures/opt/options.json

@@ -0,0 +1,3 @@
+{
+  "plugins": ["transform-es2015-for-of", "transform-flow-strip-types"]
+}

packages/babel-plugin-transform-es2015-for-of/test/fixtures/opt/typeannotation/actual.js

@@ -0,0 +1,4 @@
+for (const y of (b: Array<any>)) {}
+function a(b: Array<any>) {
+  for (const y of b) {}
+}

packages/babel-plugin-transform-es2015-for-of/test/fixtures/opt/typeannotation/expected.js

@@ -0,0 +1,28 @@
+var _arr = b;
+for (var _i = 0; _i < _arr.length; _i++) {
+  const y = _arr[_i];
+}
+function a(b) {
+  var _iteratorNormalCompletion = true;
+  var _didIteratorError = false;
+  var _iteratorError = undefined;
+
+  try {
+    for (var _iterator = b[Symbol.iterator](), _step; !(_iteratorNormalCompletion = (_step = _iterator.next()).done); _iteratorNormalCompletion = true) {
+      const y = _step.value;
+    }
+  } catch (err) {
+    _didIteratorError = true;
+    _iteratorError = err;
+  } finally {
+    try {
+      if (!_iteratorNormalCompletion && _iterator.return) {
+        _iterator.return();
+      }
+    } finally {
+      if (_didIteratorError) {
+        throw _iteratorError;
+      }
+    }
+  }
+}