From afe2f5c810b10065ba8600ae367884e59c135465 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Wed, 26 Apr 2023 18:16:28 -0700
Subject: [PATCH] Support useMemo w named functions

Supports useMemo invocations where the function is not an inline function, such
as `useMemo(someFunction, [])`.
---
 .../forget/src/Inference/InlineUseMemo.ts     | 10 +++----
 .../compiler/useMemo-named-function.expect.md | 30 +++++++++++++++++++
 .../compiler/useMemo-named-function.js        |  4 +++
 3 files changed, 38 insertions(+), 6 deletions(-)
 create mode 100644 compiler/forget/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md
 create mode 100644 compiler/forget/src/__tests__/fixtures/compiler/useMemo-named-function.js

diff --git a/compiler/forget/src/Inference/InlineUseMemo.ts b/compiler/forget/src/Inference/InlineUseMemo.ts
index ad0d78943d..8664111f87 100644
--- a/compiler/forget/src/Inference/InlineUseMemo.ts
+++ b/compiler/forget/src/Inference/InlineUseMemo.ts
@@ -17,12 +17,12 @@ import {
   GotoVariant,
   HIR,
   HIRFunction,
-  IdentifierId,
   Identifier,
+  IdentifierId,
   InstructionKind,
+  Place,
   makeInstructionId,
   makeType,
-  Place,
   reversePostorderBlocks,
   shrink,
 } from "../HIR";
@@ -97,10 +97,8 @@ export function inlineUseMemo(fn: HIRFunction): void {
             }
             const body = functions.get(lambda.identifier.id);
             if (body === undefined) {
-              CompilerError.invariant(
-                "Expected first argument to useMemo() to be a function expression",
-                lambda.loc
-              );
+              // Allow passing a named function to useMemo, eg `useMemo(someImportedFunction, [])`
+              continue;
             }
 
             if (body.loweredFunc.params.length > 0) {
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md
new file mode 100644
index 0000000000..223c3fe87e
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-named-function.expect.md
@@ -0,0 +1,30 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const x = useMemo(someHelper, []);
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = someHelper();
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const x = t0;
+  return x;
+}
+
+```
+      
\ No newline at end of file
diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useMemo-named-function.js b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-named-function.js
new file mode 100644
index 0000000000..0a9e80fec2
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/compiler/useMemo-named-function.js
@@ -0,0 +1,4 @@
+function Component(props) {
+  const x = useMemo(someHelper, []);
+  return x;
+}