Compile args to forwardRef/memo

Completes a todo (ie fixes a silly mistake) from a PR earlier in the stack, so 
we now correctly recognize and compile arguments to `React.forwardRef()` and 
`React.memo()`.

compiler/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts

@@ -455,33 +455,43 @@ function isForwardRefCallback(path: NodePath<t.Expression>): boolean {
  */
 
 function isMemoCallback(path: NodePath<t.Expression>): boolean {
-  return !!(
+  return (
     path.parentPath.isCallExpression() &&
     path.parentPath.get("callee").isExpression() &&
     isReactFunction(path.parentPath.get("callee"), "memo")
   );
 }
 
+// Adapted from the ESLint rule at
+// https://github.com/facebook/react/blob/main/packages/eslint-plugin-react-hooks/src/RulesOfHooks.js#L90-L103
 function isReactFunctionLike(
   node: NodePath<
     t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression
   >
 ): boolean {
   const functionName = getFunctionName(node);
-  if (functionName !== null) {
-    if (!isComponentName(functionName) && !isHook(functionName)) {
+  // Check if the name is component or hook like:
+  if (
+    functionName !== null &&
+    (isComponentName(functionName) || isHook(functionName))
+  ) {
+    // As an added check we also look for hook invocations or JSX
+    return callsHooksOrCreatesJsx(node);
+  }
+  // Otherwise for function or arrow function expressions, check if they
+  // appear as the argument to React.forwardRef() or React.memo():
+  if (node.isFunctionExpression() || node.isArrowFunctionExpression()) {
+    if (isForwardRefCallback(node) || isMemoCallback(node)) {
+      // As an added check we also look for hook invocations or JSX
+      return callsHooksOrCreatesJsx(node);
+    } else {
       return false;
     }
-  } else if (
-    node.isExpression() &&
-    !isForwardRefCallback(node) &&
-    !isMemoCallback(node)
-  ) {
-    return false;
-  } else {
-    return false;
   }
+  return false;
+}
 
+function callsHooksOrCreatesJsx(node: NodePath<t.Node>): boolean {
   let invokesHooks = false;
   let createsJsx = false;
   node.traverse({

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/infer-function-React-memo.expect.md

@@ -12,9 +12,17 @@ React.memo((props) => {
 ## Code
 
 ```javascript
-// @compilationMode(infer)
+import { unstable_useMemoCache as useMemoCache } from "react"; // @compilationMode(infer)
 React.memo((props) => {
-  return <div />;
+  const $ = useMemoCache(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <div />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
 });
 
 ```

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/infer-function-expression-React-memo-gating.expect.md

@@ -0,0 +1,38 @@
+
+## Input
+
+```javascript
+// @gating @compilationMode(infer)
+import React from "react";
+export default React.forwardRef(function notNamedLikeAComponent(props) {
+  return <div />;
+});
+
+```
+
+## Code
+
+```javascript
+import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
+import { unstable_useMemoCache as useMemoCache } from "react"; // @gating @compilationMode(infer)
+import React from "react";
+export default React.forwardRef(
+  isForgetEnabled_Fixtures()
+    ? function notNamedLikeAComponent(props) {
+        const $ = useMemoCache(1);
+        let t0;
+        if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+          t0 = <div />;
+          $[0] = t0;
+        } else {
+          t0 = $[0];
+        }
+        return t0;
+      }
+    : function notNamedLikeAComponent(props) {
+        return <div />;
+      }
+);
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/infer-function-expression-React-memo-gating.js

@@ -0,0 +1,5 @@
+// @gating @compilationMode(infer)
+import React from "react";
+export default React.forwardRef(function notNamedLikeAComponent(props) {
+  return <div />;
+});

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/infer-function-forwardRef.expect.md

@@ -12,9 +12,17 @@ React.forwardRef((props) => {
 ## Code
 
 ```javascript
-// @compilationMode(infer)
+import { unstable_useMemoCache as useMemoCache } from "react"; // @compilationMode(infer)
 React.forwardRef((props) => {
-  return <div />;
+  const $ = useMemoCache(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <div />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
 });
 
 ```

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/infer-skip-components-without-hooks-or-jsx.expect.md

@@ -0,0 +1,25 @@
+
+## Input
+
+```javascript
+// @compilationMode(infer)
+// This component is skipped bc it doesn't call any hooks or
+// use JSX:
+function Component(props) {
+  return render();
+}
+
+```
+
+## Code
+
+```javascript
+// @compilationMode(infer)
+// This component is skipped bc it doesn't call any hooks or
+// use JSX:
+function Component(props) {
+  return render();
+}
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/infer-skip-components-without-hooks-or-jsx.js

@@ -0,0 +1,6 @@
+// @compilationMode(infer)
+// This component is skipped bc it doesn't call any hooks or
+// use JSX:
+function Component(props) {
+  return render();
+}

compiler/packages/sprout/src/SproutTodoFilter.ts

@@ -448,6 +448,8 @@ const skipFilter = new Set([
   "infer-functions-hook-with-hook-call",
   "infer-functions-hook-with-jsx",
   "infer-function-expression-component",
+  "infer-function-expression-React-memo-gating",
+  "infer-skip-components-without-hooks-or-jsx",
 ]);
 
 export default skipFilter;