[compiler] Allow functions containing refs to be returned

Summary:
We previously were excessively strict about preventing functions that access refs from being returned--doing so is potentially valid for hooks, because the return value may only be used in an event or effect.

ghstack-source-id: cfa8bb1b54e8eb365f2de50d051bd09e09162d7b
Pull Request resolved: #30724

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts

@@ -59,20 +59,7 @@ function validateNoRefAccessInRenderImpl(
         case 'JsxExpression':
         case 'JsxFragment': {
           for (const operand of eachInstructionValueOperand(instr.value)) {
-            if (isRefValueType(operand.identifier)) {
-              errors.push({
-                severity: ErrorSeverity.InvalidReact,
-                reason:
-                  'Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef)',
-                loc: lookupLocations.get(operand.identifier.id) ?? operand.loc,
-                description:
-                  operand.identifier.name !== null &&
-                  operand.identifier.name.kind === 'named'
-                    ? `Cannot access ref value \`${operand.identifier.name.value}\``
-                    : null,
-                suggestions: null,
-              });
-            }
+            validateNoDirectRefValueAccess(errors, operand, lookupLocations);
           }
           break;
         }
@@ -232,12 +219,17 @@ function validateNoRefAccessInRenderImpl(
       }
     }
     for (const operand of eachTerminalOperand(block.terminal)) {
-      validateNoRefValueAccess(
-        errors,
-        refAccessingFunctions,
-        lookupLocations,
-        operand,
-      );
+      if (block.terminal.kind !== 'return') {
+        validateNoRefValueAccess(
+          errors,
+          refAccessingFunctions,
+          lookupLocations,
+          operand,
+        );
+      } else {
+        // Allow functions containing refs to be returned, but not direct ref values
+        validateNoDirectRefValueAccess(errors, operand, lookupLocations);
+      }
     }
   }
 
@@ -297,3 +289,24 @@ function validateNoRefAccess(
     });
   }
 }
+
+function validateNoDirectRefValueAccess(
+  errors: CompilerError,
+  operand: Place,
+  lookupLocations: Map<IdentifierId, SourceLocation>,
+): void {
+  if (isRefValueType(operand.identifier)) {
+    errors.push({
+      severity: ErrorSeverity.InvalidReact,
+      reason:
+        'Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef)',
+      loc: lookupLocations.get(operand.identifier.id) ?? operand.loc,
+      description:
+        operand.identifier.name !== null &&
+        operand.identifier.name.kind === 'named'
+          ? `Cannot access ref value \`${operand.identifier.name.value}\``
+          : null,
+      suggestions: null,
+    });
+  }
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/return-ref-callback.expect.md

@@ -0,0 +1,55 @@
+
+## Input
+
+```javascript
+// @flow @validateRefAccessDuringRender @validatePreserveExistingMemoizationGuarantees
+
+import {useRef} from 'react';
+
+component Foo() {
+  const ref = useRef();
+
+  const s = () => {
+    return ref.current;
+  };
+
+  return s;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+
+import { useRef } from "react";
+
+function Foo() {
+  const $ = _c(1);
+  const ref = useRef();
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => ref.current;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const s = t0;
+  return s;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) "[[ function params=0 ]]"

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.return-ref-callback.js → compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/return-ref-callback.js

@@ -1,5 +1,7 @@
 // @flow @validateRefAccessDuringRender @validatePreserveExistingMemoizationGuarantees
 
+import {useRef} from 'react';
+
 component Foo() {
   const ref = useRef();