From e502e69b2dedb5bd6bd08ef7d70a2b853294d953 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Mon, 6 Nov 2023 08:33:30 -0800
Subject: [PATCH] Restore React.useMemo validation

This got lost by moving the validation earlier.
---
 .../src/Validation/ValidateUseMemo.ts         | 11 ++++++++++
 ...alid-ReactUseMemo-async-callback.expect.md | 21 +++++++++++++++++++
 ...ror.invalid-ReactUseMemo-async-callback.js |  6 ++++++
 3 files changed, 38 insertions(+)
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.invalid-ReactUseMemo-async-callback.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.invalid-ReactUseMemo-async-callback.js

diff --git a/compiler/packages/babel-plugin-react-forget/src/Validation/ValidateUseMemo.ts b/compiler/packages/babel-plugin-react-forget/src/Validation/ValidateUseMemo.ts
index b4644398ef..812ce08708 100644
--- a/compiler/packages/babel-plugin-react-forget/src/Validation/ValidateUseMemo.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/Validation/ValidateUseMemo.ts
@@ -10,6 +10,7 @@ import { FunctionExpression, HIRFunction, IdentifierId } from "../HIR";
 
 export function validateUseMemo(fn: HIRFunction): void {
   const useMemos = new Set<IdentifierId>();
+  const react = new Set<IdentifierId>();
   const functions = new Map<IdentifierId, FunctionExpression>();
   for (const [, block] of fn.body.blocks) {
     for (const { lvalue, value } of block.instructions) {
@@ -17,6 +18,16 @@ export function validateUseMemo(fn: HIRFunction): void {
         case "LoadGlobal": {
           if (value.name === "useMemo") {
             useMemos.add(lvalue.identifier.id);
+          } else if (value.name === "React") {
+            react.add(lvalue.identifier.id);
+          }
+          break;
+        }
+        case "PropertyLoad": {
+          if (react.has(value.object.identifier.id)) {
+            if (value.property === "useMemo") {
+              useMemos.add(lvalue.identifier.id);
+            }
           }
           break;
         }
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.invalid-ReactUseMemo-async-callback.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.invalid-ReactUseMemo-async-callback.expect.md
new file mode 100644
index 0000000000..daf508978e
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.invalid-ReactUseMemo-async-callback.expect.md
@@ -0,0 +1,21 @@
+
+## Input
+
+```javascript
+function component(a, b) {
+  let x = React.useMemo(async () => {
+    await a;
+  }, []);
+  return x;
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidReact: useMemo callbacks may not be async or generator functions (2:4)
+```
+          
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.invalid-ReactUseMemo-async-callback.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.invalid-ReactUseMemo-async-callback.js
new file mode 100644
index 0000000000..e5a7eb0cca
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.invalid-ReactUseMemo-async-callback.js
@@ -0,0 +1,6 @@
+function component(a, b) {
+  let x = React.useMemo(async () => {
+    await a;
+  }, []);
+  return x;
+}