Fix a few type inferences (#5835)

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

@@ -1,9 +1,4 @@
 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;
   `);
@@ -100,44 +95,21 @@ export default function({ messages, template, types: t }) {
   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 (optimize(path, path.get("right"))) {
-          return;
+        const right = path.get("right");
+        if (
+          right.isArrayExpression() ||
+          right.isGenericType("Array") ||
+          t.isArrayTypeAnnotation(right.getTypeAnnotation())
+        ) {
+          return replaceWithArray(path);
         }
 
         let callback = spec;

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

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

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

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

packages/babel-traverse/src/path/inference/inferer-reference.js

@@ -11,7 +11,11 @@ export default function(node: Object) {
     if (binding.identifier.typeAnnotation) {
       return binding.identifier.typeAnnotation;
     } else {
-      return getTypeAnnotationBindingConstantViolations(this, node.name);
+      return getTypeAnnotationBindingConstantViolations(
+        binding,
+        this,
+        node.name,
+      );
     }
   }
 
@@ -25,11 +29,8 @@ export default function(node: Object) {
   }
 }
 
-function getTypeAnnotationBindingConstantViolations(path, name) {
-  const binding = path.scope.getBinding(name);
-
+function getTypeAnnotationBindingConstantViolations(binding, path, name) {
   const types = [];
-  path.typeAnnotation = t.unionTypeAnnotation(types);
 
   const functionConstantViolations = [];
   let constantViolations = getConstantViolationsBefore(
@@ -38,7 +39,7 @@ function getTypeAnnotationBindingConstantViolations(path, name) {
     functionConstantViolations,
   );
 
-  const testType = getConditionalAnnotation(path, name);
+  const testType = getConditionalAnnotation(binding, path, name);
   if (testType) {
     const testConstantViolations = getConstantViolationsBefore(
       binding,
@@ -118,18 +119,20 @@ function inferAnnotationFromBinaryExpression(name, path) {
   } else if (right.isIdentifier({ name })) {
     target = left;
   }
+
   if (target) {
     if (operator === "===") {
       return target.getTypeAnnotation();
-    } else if (t.BOOLEAN_NUMBER_BINARY_OPERATORS.indexOf(operator) >= 0) {
-      return t.numberTypeAnnotation();
-    } else {
-      return;
     }
-  } else {
-    if (operator !== "===") return;
+    if (t.BOOLEAN_NUMBER_BINARY_OPERATORS.indexOf(operator) >= 0) {
+      return t.numberTypeAnnotation();
+    }
+
+    return;
   }
 
+  if (operator !== "===" && operator !== "==") return;
+
   //
   let typeofPath;
   let typePath;
@@ -140,7 +143,10 @@ function inferAnnotationFromBinaryExpression(name, path) {
     typeofPath = right;
     typePath = left;
   }
-  if (!typePath && !typeofPath) return;
+
+  if (!typeofPath) return;
+  // and that the argument of the typeof path references us!
+  if (!typeofPath.get("argument").isIdentifier({ name })) return;
 
   // ensure that the type path is a Literal
   typePath = typePath.resolve();
@@ -150,56 +156,56 @@ function inferAnnotationFromBinaryExpression(name, path) {
   const typeValue = typePath.node.value;
   if (typeof typeValue !== "string") return;
 
-  // and that the argument of the typeof path references us!
-  if (!typeofPath.get("argument").isIdentifier({ name })) return;
-
   // turn type value into a type annotation
-  return t.createTypeAnnotationBasedOnTypeof(typePath.node.value);
+  return t.createTypeAnnotationBasedOnTypeof(typeValue);
 }
 
-function getParentConditionalPath(path) {
+function getParentConditionalPath(binding, path, name) {
   let parentPath;
   while ((parentPath = path.parentPath)) {
     if (parentPath.isIfStatement() || parentPath.isConditionalExpression()) {
       if (path.key === "test") {
         return;
-      } else {
-        return parentPath;
       }
-    } else {
-      path = parentPath;
+
+      return parentPath;
     }
+    if (parentPath.isFunction()) {
+      if (parentPath.parentPath.scope.getBinding(name) !== binding) return;
+    }
+
+    path = parentPath;
   }
 }
 
-function getConditionalAnnotation(path, name) {
-  const ifStatement = getParentConditionalPath(path);
+function getConditionalAnnotation(binding, path, name) {
+  const ifStatement = getParentConditionalPath(binding, path, name);
   if (!ifStatement) return;
 
   const test = ifStatement.get("test");
   const paths = [test];
   const types = [];
 
-  do {
-    const path = paths.shift().resolve();
+  for (let i = 0; i < paths.length; i++) {
+    const path = paths[i];
 
     if (path.isLogicalExpression()) {
-      paths.push(path.get("left"));
-      paths.push(path.get("right"));
-    }
-
-    if (path.isBinaryExpression()) {
+      if (path.node.operator === "&&") {
+        paths.push(path.get("left"));
+        paths.push(path.get("right"));
+      }
+    } else if (path.isBinaryExpression()) {
       const type = inferAnnotationFromBinaryExpression(name, path);
       if (type) types.push(type);
     }
-  } while (paths.length);
+  }
 
   if (types.length) {
     return {
       typeAnnotation: t.createUnionTypeAnnotation(types),
       ifStatement,
     };
-  } else {
-    return getConditionalAnnotation(ifStatement, name);
   }
+
+  return getConditionalAnnotation(ifStatement, name);
 }

packages/babel-traverse/src/path/inference/inferers.js

@@ -145,7 +145,22 @@ export {
   Func as ClassDeclaration,
 };
 
+const isArrayFrom = t.buildMatchMemberExpression("Array.from");
+const isObjectKeys = t.buildMatchMemberExpression("Object.keys");
+const isObjectValues = t.buildMatchMemberExpression("Object.values");
+const isObjectEntries = t.buildMatchMemberExpression("Object.entries");
 export function CallExpression() {
+  const { callee } = this.node;
+  if (isObjectKeys(callee)) {
+    return t.arrayTypeAnnotation(t.stringTypeAnnotation());
+  } else if (isArrayFrom(callee) || isObjectValues(callee)) {
+    return t.arrayTypeAnnotation(t.anyTypeAnnotation());
+  } else if (isObjectEntries(callee)) {
+    return t.arrayTypeAnnotation(
+      t.tupleTypeAnnotation([t.stringTypeAnnotation(), t.anyTypeAnnotation()]),
+    );
+  }
+
   return resolveCall(this.get("callee"));
 }
 

packages/babel-traverse/test/inference.js

@@ -259,5 +259,95 @@ describe("inference", function() {
         "should be RegExp",
       );
     });
+    it("should infer constant identifier", function() {
+      const path = getPath("const x = 0; x").get("body.1.expression");
+      const type = path.getTypeAnnotation();
+      assert.ok(t.isNumberTypeAnnotation(type), "should be number");
+    });
+    it("should infer indirect constant identifier", function() {
+      const path = getPath("const x = 0; const y = x; y").get(
+        "body.2.expression",
+      );
+      const type = path.getTypeAnnotation();
+      assert.ok(t.isNumberTypeAnnotation(type), "should be number");
+    });
+    it("should infer identifier type from if statement (===)", function() {
+      const path = getPath(
+        `function test(x) {
+        if (x === true) x;
+      }`,
+      ).get("body.0.body.body.0.consequent.expression");
+      const type = path.getTypeAnnotation();
+      assert.ok(t.isBooleanTypeAnnotation(type), "should be boolean");
+    });
+    it("should infer identifier type from if statement (typeof)", function() {
+      let path = getPath(
+        `function test(x) {
+        if (typeof x == 'string') x;
+      }`,
+      ).get("body.0.body.body.0.consequent.expression");
+      let type = path.getTypeAnnotation();
+      assert.ok(t.isStringTypeAnnotation(type), "should be string");
+      path = getPath(
+        `function test(x) {
+        if (typeof x === 'number') x;
+      }`,
+      ).get("body.0.body.body.0.consequent.expression");
+      type = path.getTypeAnnotation();
+      assert.ok(t.isNumberTypeAnnotation(type), "should be string");
+    });
+    it("should infer identifier type from if statement (&&)", function() {
+      let path = getPath(
+        `function test(x) {
+        if (typeof x == 'string' && x === 3) x;
+      }`,
+      ).get("body.0.body.body.0.consequent.expression");
+      let type = path.getTypeAnnotation();
+      assert.ok(t.isUnionTypeAnnotation(type), "should be a union");
+      assert.ok(
+        t.isStringTypeAnnotation(type.types[0]),
+        "first type in union should be string",
+      );
+      assert.ok(
+        t.isNumberTypeAnnotation(type.types[1]),
+        "second type in union should be number",
+      );
+      path = getPath(
+        `function test(x) {
+        if (true && x === 3) x;
+      }`,
+      ).get("body.0.body.body.0.consequent.expression");
+      type = path.getTypeAnnotation();
+      assert.ok(t.isNumberTypeAnnotation(type), "should be number");
+      path = getPath(
+        `function test(x) {
+        if (x === 'test' && true) x;
+      }`,
+      ).get("body.0.body.body.0.consequent.expression");
+      type = path.getTypeAnnotation();
+      assert.ok(t.isStringTypeAnnotation(type), "should be string");
+    });
+    it("should infer identifier type from if statement (||)", function() {
+      const path = getPath(
+        `function test(x) {
+        if (typeof x == 'string' || x === 3) x;
+      }`,
+      ).get("body.0.body.body.0.consequent.expression");
+      const type = path.getTypeAnnotation();
+      assert.ok(t.isAnyTypeAnnotation(type), "should be a any type");
+    });
+    it("should not infer identifier type from incorrect binding", function() {
+      const path = getPath(
+        `function outer(x) {
+        if (x === 3) {
+          function inner(x) {
+            x;
+          }
+        }
+      }`,
+      ).get("body.0.body.body.0.consequent.body.0.body.body.0.expression");
+      const type = path.getTypeAnnotation();
+      assert.ok(t.isAnyTypeAnnotation(type), "should be a any type");
+    });
   });
 });