From 6cc29209661e30746b2c77923f2eb0fdb402a327 Mon Sep 17 00:00:00 2001
From: Federico Ciardi <fed.ciardi@gmail.com>
Date: Sun, 31 Jan 2021 20:39:25 +0100
Subject: [PATCH 1/6] fix: infer array type from outside of arrow function

---
 .../fixtures/spec/function-parent/input.js    | 14 +++++++++++++
 .../fixtures/spec/function-parent/output.js   | 21 +++++++++++++++++++
 .../test/fixtures/spec/infer-type/input.js    |  3 +++
 .../test/fixtures/spec/infer-type/output.js   |  5 +++++
 .../babel-traverse/src/path/introspection.ts  | 12 ++++++++---
 5 files changed, 52 insertions(+), 3 deletions(-)
 create mode 100644 packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
 create mode 100644 packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
 create mode 100644 packages/babel-plugin-transform-for-of/test/fixtures/spec/infer-type/input.js
 create mode 100644 packages/babel-plugin-transform-for-of/test/fixtures/spec/infer-type/output.js

diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
new file mode 100644
index 000000000000..9e45f2a92b3e
--- /dev/null
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
@@ -0,0 +1,14 @@
+const keys = []
+function a() {
+  for (const key of keys) {}
+}
+
+const b = () => {
+  for (const key of keys) {}
+}
+
+const c = () => {
+  const d = () => {
+    for (const key of keys) {}
+  }
+}
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
new file mode 100644
index 000000000000..3aca7fca4865
--- /dev/null
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
@@ -0,0 +1,21 @@
+const keys = [];
+
+function a() {
+  for (var _i = 0, _keys = keys; _i < _keys.length; _i++) {
+    const key = _keys[_i];
+  }
+}
+
+const b = () => {
+  for (var _i2 = 0, _keys2 = keys; _i2 < _keys2.length; _i2++) {
+    const key = _keys2[_i2];
+  }
+};
+
+const c = () => {
+  const d = () => {
+    for (var _i3 = 0, _keys3 = keys; _i3 < _keys3.length; _i3++) {
+      const key = _keys3[_i3];
+    }
+  };
+};
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/infer-type/input.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/infer-type/input.js
new file mode 100644
index 000000000000..9d91ad61c691
--- /dev/null
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/infer-type/input.js
@@ -0,0 +1,3 @@
+const foo = []
+
+for (const i of foo) {}
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/infer-type/output.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/infer-type/output.js
new file mode 100644
index 000000000000..5a965d2c1f83
--- /dev/null
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/infer-type/output.js
@@ -0,0 +1,5 @@
+const foo = [];
+
+for (var _i = 0, _foo = foo; _i < _foo.length; _i++) {
+  const i = _foo[_i];
+}
diff --git a/packages/babel-traverse/src/path/introspection.ts b/packages/babel-traverse/src/path/introspection.ts
index 0e01e4f20cc9..b7782eacf468 100644
--- a/packages/babel-traverse/src/path/introspection.ts
+++ b/packages/babel-traverse/src/path/introspection.ts
@@ -219,7 +219,7 @@ export function willIMaybeExecuteBefore(this: NodePath, target): boolean {
   return this._guessExecutionStatusRelativeTo(target) !== "after";
 }
 
-function getOuterFunction(path) {
+function getOuterFunction(path: NodePath) {
   return (path.scope.getFunctionParent() || path.scope.getProgramParent()).path;
 }
 
@@ -385,8 +385,11 @@ export function _guessExecutionStatusRelativeToDifferentFunctions(
   this: NodePath,
   target: NodePath,
 ): RelativeExecutionStatus {
+  const targetIsArrowExpressionDeclaration =
+    target.isArrowFunctionExpression() &&
+    t.isVariableDeclarator(target.container);
   if (
-    !target.isFunctionDeclaration() ||
+    (!target.isFunctionDeclaration() && !targetIsArrowExpressionDeclaration) ||
     target.parentPath.isExportDeclaration()
   ) {
     return "unknown";
@@ -396,7 +399,10 @@ export function _guessExecutionStatusRelativeToDifferentFunctions(
   // then we can be a bit smarter and handle cases where the function is either
   // a. not called at all (part of an export)
   // b. called directly
-  const binding = target.scope.getBinding(target.node.id.name);
+  const bindingName = targetIsArrowExpressionDeclaration
+    ? target.container.id.name
+    : target.node.id.name;
+  const binding = target.scope.getBinding(bindingName);
 
   // no references!
   if (!binding.references) return "before";

From d8f7c905c5ddeed4f96b8564707cf56f1d45536f Mon Sep 17 00:00:00 2001
From: Federico Ciardi <fed.ciardi@gmail.com>
Date: Mon, 1 Feb 2021 12:38:00 +0100
Subject: [PATCH 2/6] fix: apply code review suggestions

---
 .../fixtures/spec/function-parent/input.js    |  8 +++--
 .../fixtures/spec/function-parent/output.js   | 14 ++++++---
 .../babel-traverse/src/path/introspection.ts  | 14 ++++-----
 packages/babel-traverse/test/inference.js     | 30 +++++++++++++++++++
 4 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
index 9e45f2a92b3e..19b88de9e868 100644
--- a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
@@ -7,8 +7,12 @@ const b = () => {
   for (const key of keys) {}
 }
 
-const c = () => {
-  const d = () => {
+const c = function () {
+  for (const key of keys) {}
+}
+
+const d = () => {
+  const e = () => {
     for (const key of keys) {}
   }
 }
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
index 3aca7fca4865..34426431bf28 100644
--- a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
@@ -12,10 +12,16 @@ const b = () => {
   }
 };
 
-const c = () => {
-  const d = () => {
-    for (var _i3 = 0, _keys3 = keys; _i3 < _keys3.length; _i3++) {
-      const key = _keys3[_i3];
+const c = function () {
+  for (var _i3 = 0, _keys3 = keys; _i3 < _keys3.length; _i3++) {
+    const key = _keys3[_i3];
+  }
+};
+
+const d = () => {
+  const e = () => {
+    for (var _i4 = 0, _keys4 = keys; _i4 < _keys4.length; _i4++) {
+      const key = _keys4[_i4];
     }
   };
 };
diff --git a/packages/babel-traverse/src/path/introspection.ts b/packages/babel-traverse/src/path/introspection.ts
index b7782eacf468..85e61564c30b 100644
--- a/packages/babel-traverse/src/path/introspection.ts
+++ b/packages/babel-traverse/src/path/introspection.ts
@@ -385,11 +385,11 @@ export function _guessExecutionStatusRelativeToDifferentFunctions(
   this: NodePath,
   target: NodePath,
 ): RelativeExecutionStatus {
-  const targetIsArrowExpressionDeclaration =
-    target.isArrowFunctionExpression() &&
-    t.isVariableDeclarator(target.container);
+  const targetIsGenericFunctionExpression =
+    (target.isArrowFunctionExpression() || target.isFunctionExpression()) &&
+    t.isVariableDeclarator(target.parent);
   if (
-    (!target.isFunctionDeclaration() && !targetIsArrowExpressionDeclaration) ||
+    (!target.isFunctionDeclaration() && !targetIsGenericFunctionExpression) ||
     target.parentPath.isExportDeclaration()
   ) {
     return "unknown";
@@ -399,9 +399,9 @@ export function _guessExecutionStatusRelativeToDifferentFunctions(
   // then we can be a bit smarter and handle cases where the function is either
   // a. not called at all (part of an export)
   // b. called directly
-  const bindingName = targetIsArrowExpressionDeclaration
-    ? target.container.id.name
-    : target.node.id.name;
+  const bindingName = targetIsGenericFunctionExpression
+    ? target.parent["id"]["name"]
+    : target.node["id"]["name"];
   const binding = target.scope.getBinding(bindingName);
 
   // no references!
diff --git a/packages/babel-traverse/test/inference.js b/packages/babel-traverse/test/inference.js
index dde8bd22746e..b86c45660991 100644
--- a/packages/babel-traverse/test/inference.js
+++ b/packages/babel-traverse/test/inference.js
@@ -149,6 +149,36 @@ describe("inference", function () {
         t.isGenericTypeAnnotation(type) && type.id.name === "Function",
       ).toBeTruthy();
     });
+    it("should infer Array type from outside of function declaration", function () {
+      const path = getPath(`
+      const arr = []; function func() { const a = arr }`).get(
+        "body.1.body.body.0.declarations.0",
+      );
+      const type = path.getTypeAnnotation();
+      expect(
+        t.isGenericTypeAnnotation(type) && type.id.name === "Array",
+      ).toBeTruthy();
+    });
+    it("should infer Array type from outside of function expression", function () {
+      const path = getPath(`
+        const arr = []; const func = function () { const a = arr }`).get(
+        "body.1.declarations.0.init.body.body.0.declarations.0",
+      );
+      const type = path.getTypeAnnotation();
+      expect(
+        t.isGenericTypeAnnotation(type) && type.id.name === "Array",
+      ).toBeTruthy();
+    });
+    it("should infer Array type from outside of arrow function", function () {
+      const path = getPath(`
+        const arr = []; const func = () => { const a = arr }`).get(
+        "body.1.declarations.0.init.body.body.0.declarations.0",
+      );
+      const type = path.getTypeAnnotation();
+      expect(
+        t.isGenericTypeAnnotation(type) && type.id.name === "Array",
+      ).toBeTruthy();
+    });
     it("should infer call return type using function", function () {
       const path = getPath("(function (): string {})()")
         .get("body")[0]

From 3fe8a8c4cfed5ccec4bbaf65535700e426ccd056 Mon Sep 17 00:00:00 2001
From: Federico Ciardi <fed.ciardi@gmail.com>
Date: Tue, 2 Feb 2021 08:40:38 +0100
Subject: [PATCH 3/6] Improve type casting and bail out when non-identifier
 declarator

---
 .../test/fixtures/spec/function-parent/input.js |  4 ++++
 .../fixtures/spec/function-parent/output.js     | 17 +++++++++++++++++
 .../babel-traverse/src/path/introspection.ts    |  7 ++++---
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
index 19b88de9e868..22e8dbce861e 100644
--- a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
@@ -16,3 +16,7 @@ const d = () => {
     for (const key of keys) {}
   }
 }
+
+const { foo } = () => {
+  for (const key of keys) {}
+}
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
index 34426431bf28..a178832e6e6c 100644
--- a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
@@ -25,3 +25,20 @@ const d = () => {
     }
   };
 };
+
+const {
+  foo
+} = () => {
+  var _iterator = babelHelpers.createForOfIteratorHelper(keys),
+      _step;
+
+  try {
+    for (_iterator.s(); !(_step = _iterator.n()).done;) {
+      const key = _step.value;
+    }
+  } catch (err) {
+    _iterator.e(err);
+  } finally {
+    _iterator.f();
+  }
+};
diff --git a/packages/babel-traverse/src/path/introspection.ts b/packages/babel-traverse/src/path/introspection.ts
index 85e61564c30b..a3aa04131bfb 100644
--- a/packages/babel-traverse/src/path/introspection.ts
+++ b/packages/babel-traverse/src/path/introspection.ts
@@ -387,7 +387,8 @@ export function _guessExecutionStatusRelativeToDifferentFunctions(
 ): RelativeExecutionStatus {
   const targetIsGenericFunctionExpression =
     (target.isArrowFunctionExpression() || target.isFunctionExpression()) &&
-    t.isVariableDeclarator(target.parent);
+    t.isVariableDeclarator(target.parent) &&
+    t.isIdentifier(target.parent.id);
   if (
     (!target.isFunctionDeclaration() && !targetIsGenericFunctionExpression) ||
     target.parentPath.isExportDeclaration()
@@ -400,8 +401,8 @@ export function _guessExecutionStatusRelativeToDifferentFunctions(
   // a. not called at all (part of an export)
   // b. called directly
   const bindingName = targetIsGenericFunctionExpression
-    ? target.parent["id"]["name"]
-    : target.node["id"]["name"];
+    ? (target.parent as t.VariableDeclarator).id.name
+    : (target.node as t.FunctionDeclaration).id.name;
   const binding = target.scope.getBinding(bindingName);
 
   // no references!

From aa602c95b486b33d10459f793e46717680f8710c Mon Sep 17 00:00:00 2001
From: Federico Ciardi <fed.ciardi@gmail.com>
Date: Tue, 2 Feb 2021 20:45:02 +0100
Subject: [PATCH 4/6] fix: type casting

---
 packages/babel-traverse/src/path/introspection.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/babel-traverse/src/path/introspection.ts b/packages/babel-traverse/src/path/introspection.ts
index a3aa04131bfb..dc1a82333b37 100644
--- a/packages/babel-traverse/src/path/introspection.ts
+++ b/packages/babel-traverse/src/path/introspection.ts
@@ -401,7 +401,7 @@ export function _guessExecutionStatusRelativeToDifferentFunctions(
   // a. not called at all (part of an export)
   // b. called directly
   const bindingName = targetIsGenericFunctionExpression
-    ? (target.parent as t.VariableDeclarator).id.name
+    ? ((target.parent as t.VariableDeclarator).id as t.Identifier).name
     : (target.node as t.FunctionDeclaration).id.name;
   const binding = target.scope.getBinding(bindingName);
 

From 3dc7a6e5ee4c3351a5bd6aa5f39f77664d259eb9 Mon Sep 17 00:00:00 2001
From: Federico Ciardi <fed.ciardi@gmail.com>
Date: Sat, 20 Feb 2021 15:59:01 +0100
Subject: [PATCH 5/6] feat: guess execution time when function is returned from
 another function

---
 packages/babel-traverse/src/path/introspection.ts | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/packages/babel-traverse/src/path/introspection.ts b/packages/babel-traverse/src/path/introspection.ts
index dc1a82333b37..efd8402a17cf 100644
--- a/packages/babel-traverse/src/path/introspection.ts
+++ b/packages/babel-traverse/src/path/introspection.ts
@@ -419,7 +419,10 @@ export function _guessExecutionStatusRelativeToDifferentFunctions(
     const childOfFunction = !!path.find(path => path.node === target.node);
     if (childOfFunction) continue;
 
-    if (path.key !== "callee" || !path.parentPath.isCallExpression()) {
+    if (
+      (path.key !== "callee" || !path.parentPath.isCallExpression()) &&
+      !path.parentPath.isReturnStatement()
+    ) {
       // This function is passed as a reference, so we don't
       // know when it will be called.
       return "unknown";

From 20fad3c5f1e9e8b74aee2f2b250d11e0495d4cdc Mon Sep 17 00:00:00 2001
From: Federico Ciardi <fed.ciardi@gmail.com>
Date: Sat, 20 Feb 2021 16:12:52 +0100
Subject: [PATCH 6/6] Update tests

---
 .../spec/exported-function-parent/input.mjs   | 10 +++++
 .../spec/exported-function-parent/output.mjs  | 22 ++++++++++
 .../fixtures/spec/function-parent/input.js    |  7 +--
 .../fixtures/spec/function-parent/output.js   |  8 ----
 .../spec/nested-function-parent/input.js      | 24 ++++++++++
 .../spec/nested-function-parent/output.js     | 44 +++++++++++++++++++
 6 files changed, 101 insertions(+), 14 deletions(-)
 create mode 100644 packages/babel-plugin-transform-for-of/test/fixtures/spec/exported-function-parent/input.mjs
 create mode 100644 packages/babel-plugin-transform-for-of/test/fixtures/spec/exported-function-parent/output.mjs
 create mode 100644 packages/babel-plugin-transform-for-of/test/fixtures/spec/nested-function-parent/input.js
 create mode 100644 packages/babel-plugin-transform-for-of/test/fixtures/spec/nested-function-parent/output.js

diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/exported-function-parent/input.mjs b/packages/babel-plugin-transform-for-of/test/fixtures/spec/exported-function-parent/input.mjs
new file mode 100644
index 000000000000..ad0bb3453076
--- /dev/null
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/exported-function-parent/input.mjs
@@ -0,0 +1,10 @@
+const keys = []
+export const foo = () => {
+  for (const key of keys) {}
+}
+
+export const bar = () => {
+  const keys = []
+  for (const key of keys) {}
+}
+
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/exported-function-parent/output.mjs b/packages/babel-plugin-transform-for-of/test/fixtures/spec/exported-function-parent/output.mjs
new file mode 100644
index 000000000000..00997d4d77c1
--- /dev/null
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/exported-function-parent/output.mjs
@@ -0,0 +1,22 @@
+const keys = [];
+export const foo = () => {
+  var _iterator = babelHelpers.createForOfIteratorHelper(keys),
+      _step;
+
+  try {
+    for (_iterator.s(); !(_step = _iterator.n()).done;) {
+      const key = _step.value;
+    }
+  } catch (err) {
+    _iterator.e(err);
+  } finally {
+    _iterator.f();
+  }
+};
+export const bar = () => {
+  const keys = [];
+
+  for (var _i = 0, _keys = keys; _i < _keys.length; _i++) {
+    const key = _keys[_i];
+  }
+};
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
index 22e8dbce861e..8f79a4beb121 100644
--- a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/input.js
@@ -1,4 +1,5 @@
 const keys = []
+
 function a() {
   for (const key of keys) {}
 }
@@ -11,12 +12,6 @@ const c = function () {
   for (const key of keys) {}
 }
 
-const d = () => {
-  const e = () => {
-    for (const key of keys) {}
-  }
-}
-
 const { foo } = () => {
   for (const key of keys) {}
 }
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
index a178832e6e6c..c42b5c524243 100644
--- a/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/function-parent/output.js
@@ -18,14 +18,6 @@ const c = function () {
   }
 };
 
-const d = () => {
-  const e = () => {
-    for (var _i4 = 0, _keys4 = keys; _i4 < _keys4.length; _i4++) {
-      const key = _keys4[_i4];
-    }
-  };
-};
-
 const {
   foo
 } = () => {
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/nested-function-parent/input.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/nested-function-parent/input.js
new file mode 100644
index 000000000000..eb6d64d0917c
--- /dev/null
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/nested-function-parent/input.js
@@ -0,0 +1,24 @@
+const keys = []
+const d = () => {
+  const e = () => {
+    for (const key of keys) {}
+  }
+}
+
+const a = () => {
+  const keys = []
+  const c = () => {
+    for (const key of keys) {}
+  }
+  return c
+}
+
+var _keys = []
+const b = () => {
+  const c = () => {
+    for (const key of _keys) {}
+  }
+  return c
+}
+
+_keys = 'foo'
diff --git a/packages/babel-plugin-transform-for-of/test/fixtures/spec/nested-function-parent/output.js b/packages/babel-plugin-transform-for-of/test/fixtures/spec/nested-function-parent/output.js
new file mode 100644
index 000000000000..7e8bfcd66241
--- /dev/null
+++ b/packages/babel-plugin-transform-for-of/test/fixtures/spec/nested-function-parent/output.js
@@ -0,0 +1,44 @@
+const keys = [];
+
+const d = () => {
+  const e = () => {
+    for (var _i = 0, _keys2 = keys; _i < _keys2.length; _i++) {
+      const key = _keys2[_i];
+    }
+  };
+};
+
+const a = () => {
+  const keys = [];
+
+  const c = () => {
+    for (var _i2 = 0, _keys3 = keys; _i2 < _keys3.length; _i2++) {
+      const key = _keys3[_i2];
+    }
+  };
+
+  return c;
+};
+
+var _keys = [];
+
+const b = () => {
+  const c = () => {
+    var _iterator = babelHelpers.createForOfIteratorHelper(_keys),
+        _step;
+
+    try {
+      for (_iterator.s(); !(_step = _iterator.n()).done;) {
+        const key = _step.value;
+      }
+    } catch (err) {
+      _iterator.e(err);
+    } finally {
+      _iterator.f();
+    }
+  };
+
+  return c;
+};
+
+_keys = 'foo';