From 5775102b852ea53c53a0d1bc3349c85a735f246e Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Thu, 8 Jun 2023 14:02:04 -0400
Subject: [PATCH] InferTypes traverses function expressions

---
 .../src/Inference/AnalyseFunctions.ts                     | 2 --
 .../src/TypeInference/InferTypes.ts                       | 8 ++++++--
 .../babel-plugin-react-forget/src/TypeInference/index.ts  | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts b/compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts
index 4f1c42d422..103b01dd05 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts
@@ -16,7 +16,6 @@ import {
   Place,
   ReactiveScopeDependency,
 } from "../HIR";
-import { inferTypes } from "../TypeInference";
 import { logHIRFunction } from "../Utils/logger";
 import { inferMutableRanges } from "./InferMutableRanges";
 import inferReferenceEffects from "./InferReferenceEffects";
@@ -87,7 +86,6 @@ export default function analyseFunctions(func: HIRFunction): void {
 }
 
 function lower(func: HIRFunction): void {
-  inferTypes(func);
   analyseFunctions(func);
   inferReferenceEffects(func, { isFunctionExpression: true });
   inferMutableRanges(func);
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
index 058b16a697..0b8acae0be 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
@@ -45,7 +45,7 @@ function isPrimitiveBinaryOp(op: t.BinaryExpression["operator"]): boolean {
   }
 }
 
-export default function (func: HIRFunction): void {
+export function inferTypes(func: HIRFunction): void {
   const unifier = new Unifier(func.env);
   for (const e of generate(func)) {
     unifier.unify(e.left, e.right);
@@ -244,6 +244,11 @@ function* generateInstructionTypes(
       break;
     }
 
+    case "FunctionExpression": {
+      inferTypes(value.loweredFunc);
+      break;
+    }
+
     case "DeclareLocal":
     case "DeclareContext":
     case "NewExpression":
@@ -253,7 +258,6 @@ function* generateInstructionTypes(
     case "PropertyStore":
     case "ComputedStore":
     case "ComputedLoad":
-    case "FunctionExpression":
     case "TaggedTemplateExpression":
     case "Await":
     case "NextIterableOf":
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/index.ts b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/index.ts
index 009c082e3c..adc3a55c7a 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/index.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/index.ts
@@ -5,4 +5,4 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-export { default as inferTypes } from "./InferTypes";
+export { inferTypes } from "./InferTypes";