fix: support some TSTypes in inferrer

packages/babel-plugin-transform-for-of/src/index.ts

@@ -207,6 +207,7 @@ export default declare((api, options: Options) => {
         if (
           right.isArrayExpression() ||
           right.isGenericType("Array") ||
+          // todo(Babel 8): the following line can be removed as it is supported by isGenericType("Array")
           t.isArrayTypeAnnotation(right.getTypeAnnotation())
         ) {
           path.replaceWith(_ForOfStatementArray(path));

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

@@ -1,3 +1,4 @@
 {
-  "plugins": ["transform-for-of", "transform-flow-strip-types"]
+  "plugins": ["transform-for-of"],
+  "presets": ["typescript"]
 }

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow-deopt/input.js

@@ -0,0 +1,4 @@
+// This won't be optimize because when `for-of` is handled, the b's type annotation has been removed by the Flow plugin
+function a(b: Array<any>) {
+  for (const y of b) {}
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow-deopt/options.json

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

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow-deopt/output.js

@@ -0,0 +1,15 @@
+// This won't be optimize because when `for-of` is handled, the b's type annotation has been removed by the Flow plugin
+function a(b) {
+  var _iterator = babelHelpers.createForOfIteratorHelper(b),
+      _step;
+
+  try {
+    for (_iterator.s(); !(_step = _iterator.n()).done;) {
+      const y = _step.value;
+    }
+  } catch (err) {
+    _iterator.e(err);
+  } finally {
+    _iterator.f();
+  }
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow/input.js

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

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow/options.json

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

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-flow/output.js

@@ -0,0 +1,3 @@
+for (var _i = 0, _arr = b; _i < _arr.length; _i++) {
+  const y = _arr[_i];
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts-deopt/input.ts

@@ -0,0 +1,4 @@
+// This won't be optimize because when `for-of` is handled, the b's type annotation has been removed by the TS plugin
+function a(b: Array<any>) {
+  for (const y of b) {}
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts-deopt/options.json

@@ -0,0 +1,4 @@
+{
+  "plugins": ["transform-for-of"],
+  "presets": ["typescript"]
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/typeannotation/output.js → packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts-deopt/output.js

@@ -1,7 +1,4 @@
-for (var _i = 0, _arr = b; _i < _arr.length; _i++) {
-  const y = _arr[_i];
-}
-
+// This won't be optimize because when `for-of` is handled, the b's type annotation has been removed by the TS plugin
 function a(b) {
   var _iterator = babelHelpers.createForOfIteratorHelper(b),
       _step;

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts/input.ts

@@ -0,0 +1 @@
+for (const y of (b as Array<any>)) {}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts/options.json

@@ -0,0 +1,4 @@
+{
+  "plugins": ["transform-for-of"],
+  "presets": ["typescript"]
+}

packages/babel-plugin-transform-for-of/test/fixtures/opt/type-annotation-array-ts/output.js

@@ -0,0 +1,3 @@
+for (var _i = 0, _arr = b; _i < _arr.length; _i++) {
+  const y = _arr[_i];
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-flow/input.ts

@@ -0,0 +1,14 @@
+function foo() {
+  const x = 1 ? a() : b();
+  const y = a() || b();
+
+  return [...x, ...y];
+}
+
+function a(): Array<number> {
+  return [];
+}
+
+function b(): Array<number> {
+  return [];
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-flow/options.json

@@ -0,0 +1,4 @@
+{
+  "plugins": ["transform-spread"],
+  "presets": ["flow"]
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-flow/output.js

@@ -0,0 +1,13 @@
+function foo() {
+  const x = 1 ? a() : b();
+  const y = a() || b();
+  return [].concat(x, y);
+}
+
+function a() {
+  return [];
+}
+
+function b() {
+  return [];
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-ts/input.ts

@@ -0,0 +1,14 @@
+function foo() {
+  const x = 1 ? a() : b();
+  const y = a() || b();
+
+  return [...x, ...y];
+}
+
+function a(): Array<number> {
+  return [];
+}
+
+function b(): Array<number> {
+  return [];
+}

packages/babel-plugin-transform-spread/test/fixtures/spread/type-reference-ts/output.js

@@ -0,0 +1,13 @@
+function foo() {
+  const x = 1 ? a() : b();
+  const y = a() || b();
+  return [].concat(x, y);
+}
+
+function a() {
+  return [];
+}
+
+function b() {
+  return [];
+}

packages/babel-traverse/src/path/inference/index.ts

@@ -3,6 +3,7 @@ import * as inferers from "./inferers";
 import {
   anyTypeAnnotation,
   isAnyTypeAnnotation,
+  isArrayTypeAnnotation,
   isBooleanTypeAnnotation,
   isEmptyTypeAnnotation,
   isFlowBaseAnnotation,
@@ -11,6 +12,10 @@ import {
   isMixedTypeAnnotation,
   isNumberTypeAnnotation,
   isStringTypeAnnotation,
+  isTSArrayType,
+  isTSTypeAnnotation,
+  isTSTypeReference,
+  isTupleTypeAnnotation,
   isTypeAnnotation,
   isUnionTypeAnnotation,
   isVoidTypeAnnotation,
@@ -29,7 +34,9 @@ export function getTypeAnnotation(this: NodePath): t.FlowType {
     return type;
   }
   type = this._getTypeAnnotation() || anyTypeAnnotation();
-  if (isTypeAnnotation(type)) type = type.typeAnnotation;
+  if (isTypeAnnotation(type) || isTSTypeAnnotation(type)) {
+    type = type.typeAnnotation;
+  }
   this.setData("typeAnnotation", type);
   return type;
 }
@@ -158,8 +165,24 @@ export function baseTypeStrictlyMatches(
 
 export function isGenericType(this: NodePath, genericName: string): boolean {
   const type = this.getTypeAnnotation();
+  if (genericName === "Array") {
+    // T[]
+    if (
+      isTSArrayType(type) ||
+      isArrayTypeAnnotation(type) ||
+      isTupleTypeAnnotation(type)
+    ) {
+      return true;
+    }
+  }
   return (
-    isGenericTypeAnnotation(type) &&
-    isIdentifier(type.id, { name: genericName })
+    (isGenericTypeAnnotation(type) &&
+      isIdentifier(type.id, {
+        name: genericName,
+      })) ||
+    (isTSTypeReference(type) &&
+      isIdentifier(type.typeName, {
+        name: genericName,
+      }))
   );
 }

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

@@ -1,17 +1,15 @@
 import type NodePath from "../index";
 import {
   BOOLEAN_NUMBER_BINARY_OPERATORS,
-  createFlowUnionType,
-  createTSUnionType,
   createTypeAnnotationBasedOnTypeof,
-  createUnionTypeAnnotation,
-  isTSTypeAnnotation,
   numberTypeAnnotation,
   voidTypeAnnotation,
 } from "@babel/types";
 import type * as t from "@babel/types";
 import type Binding from "../../scope/binding";
 
+import { createUnionType } from "./util";
+
 export default function (this: NodePath, node: any) {
   if (!this.isReferenced()) return;
 
@@ -106,15 +104,7 @@ function getTypeAnnotationBindingConstantViolations(binding, path, name) {
     return;
   }
 
-  if (isTSTypeAnnotation(types[0]) && createTSUnionType) {
-    return createTSUnionType(types);
-  }
-
-  if (createFlowUnionType) {
-    return createFlowUnionType(types);
-  }
-
-  return createUnionTypeAnnotation(types);
+  return createUnionType(types);
 }
 
 function getConstantViolationsBefore(binding: Binding, path, functions?) {
@@ -231,22 +221,8 @@ function getConditionalAnnotation<T extends t.Node>(
   }
 
   if (types.length) {
-    if (isTSTypeAnnotation(types[0]) && createTSUnionType) {
-      return {
-        typeAnnotation: createTSUnionType(types),
-        ifStatement,
-      };
-    }
-
-    if (createFlowUnionType) {
-      return {
-        typeAnnotation: createFlowUnionType(types),
-        ifStatement,
-      };
-    }
-
     return {
-      typeAnnotation: createUnionTypeAnnotation(types),
+      typeAnnotation: createUnionType(types),
       ifStatement,
     };
   }

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

@@ -8,45 +8,26 @@ import {
   arrayTypeAnnotation,
   booleanTypeAnnotation,
   buildMatchMemberExpression,
-  createFlowUnionType,
-  createTSUnionType,
-  createUnionTypeAnnotation,
   genericTypeAnnotation,
   identifier,
-  isTSTypeAnnotation,
   nullLiteralTypeAnnotation,
   numberTypeAnnotation,
   stringTypeAnnotation,
   tupleTypeAnnotation,
   unionTypeAnnotation,
   voidTypeAnnotation,
+  isIdentifier,
 } from "@babel/types";
 import type * as t from "@babel/types";
 
 export { default as Identifier } from "./inferer-reference";
 
+import { createUnionType } from "./util";
 import type NodePath from "..";
 
 export function VariableDeclarator(this: NodePath<t.VariableDeclarator>) {
-  const id = this.get("id");
-
-  if (!id.isIdentifier()) return;
-  const init = this.get("init");
-
-  let type = init.getTypeAnnotation();
-
-  if (type?.type === "AnyTypeAnnotation") {
-    // Detect "var foo = Array()" calls so we can optimize for arrays vs iterables.
-    if (
-      init.isCallExpression() &&
-      init.get("callee").isIdentifier({ name: "Array" }) &&
-      !init.scope.hasBinding("Array", true /* noGlobals */)
-    ) {
-      type = ArrayExpression();
-    }
-  }
-
-  return type;
+  if (!this.get("id").isIdentifier()) return;
+  return this.get("init").getTypeAnnotation();
 }
 
 export function TypeCastExpression(node) {
@@ -61,6 +42,10 @@ export function TSAsExpression(node) {
 
 TSAsExpression.validParent = true;
 
+export function TSNonNullExpression() {
+  return this.get("expression").getTypeAnnotation();
+}
+
 export function NewExpression(this: NodePath<t.NewExpression>, node) {
   if (this.get("callee").isIdentifier()) {
     // only resolve identifier callee
@@ -119,16 +104,7 @@ export function LogicalExpression(this: NodePath<t.LogicalExpression>) {
     this.get("right").getTypeAnnotation(),
   ];
 
-  if (isTSTypeAnnotation(argumentTypes[0]) && createTSUnionType) {
-    // @ts-expect-error Fixme: getTypeAnnotation also returns TS types
-    return createTSUnionType(argumentTypes);
-  }
-
-  if (createFlowUnionType) {
-    return createFlowUnionType(argumentTypes);
-  }
-
-  return createUnionTypeAnnotation(argumentTypes);
+  return createUnionType(argumentTypes);
 }
 
 export function ConditionalExpression(this: NodePath<t.ConditionalExpression>) {
@@ -137,16 +113,7 @@ export function ConditionalExpression(this: NodePath<t.ConditionalExpression>) {
     this.get("alternate").getTypeAnnotation(),
   ];
 
-  if (isTSTypeAnnotation(argumentTypes[0]) && createTSUnionType) {
-    // @ts-expect-error Fixme: getTypeAnnotation also returns TS types
-    return createTSUnionType(argumentTypes);
-  }
-
-  if (createFlowUnionType) {
-    return createFlowUnionType(argumentTypes);
-  }
-
-  return createUnionTypeAnnotation(argumentTypes);
+  return createUnionType(argumentTypes);
 }
 
 export function SequenceExpression(this: NodePath<t.SequenceExpression>) {
@@ -224,7 +191,12 @@ export function CallExpression(this: NodePath<t.CallExpression>) {
   const { callee } = this.node;
   if (isObjectKeys(callee)) {
     return arrayTypeAnnotation(stringTypeAnnotation());
-  } else if (isArrayFrom(callee) || isObjectValues(callee)) {
+  } else if (
+    isArrayFrom(callee) ||
+    isObjectValues(callee) ||
+    // Detect "var foo = Array()" calls so we can optimize for arrays vs iterables.
+    isIdentifier(callee, { name: "Array" })
+  ) {
     return arrayTypeAnnotation(anyTypeAnnotation());
   } else if (isObjectEntries(callee)) {
     return arrayTypeAnnotation(
@@ -245,14 +217,17 @@ function resolveCall(callee) {
   callee = callee.resolve();
 
   if (callee.isFunction()) {
-    if (callee.is("async")) {
-      if (callee.is("generator")) {
+    const { node } = callee;
+    if (node.async) {
+      if (node.generator) {
         return genericTypeAnnotation(identifier("AsyncIterator"));
       } else {
         return genericTypeAnnotation(identifier("Promise"));
       }
     } else {
-      if (callee.node.returnType) {
+      if (node.generator) {
+        return genericTypeAnnotation(identifier("Iterator"));
+      } else if (callee.node.returnType) {
         return callee.node.returnType;
       } else {
         // todo: get union type of all return arguments