[compiler] Allow all hooks to take callbacks which access refs, but ban hooks from taking direct ref value arguments

compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts

@@ -362,6 +362,17 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
       returnValueKind: ValueKind.Mutable,
     }),
   ],
+  [
+    'useImperativeHandle',
+    addHook(DEFAULT_SHAPES, {
+      positionalParams: [],
+      restParam: Effect.Freeze,
+      returnType: {kind: 'Primitive'},
+      calleeEffect: Effect.Read,
+      hookKind: 'useImperativeHandle',
+      returnValueKind: ValueKind.Frozen,
+    }),
+  ],
   [
     'useMemo',
     addHook(DEFAULT_SHAPES, {

compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts

@@ -127,6 +127,7 @@ export type HookKind =
   | 'useMemo'
   | 'useCallback'
   | 'useTransition'
+  | 'useImperativeHandle'
   | 'Custom';
 
 /*

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

@@ -22,7 +22,6 @@ import {
   eachTerminalOperand,
 } from '../HIR/visitors';
 import {Err, Ok, Result} from '../Utils/Result';
-import {isEffectHook} from './ValidateMemoizedEffectDependencies';
 
 /**
  * Validates that a function does not access a ref value during render. This includes a partial check
@@ -309,60 +308,37 @@ function validateNoRefAccessInRenderImpl(
             });
             break;
           }
-          case 'MethodCall': {
-            if (!isEffectHook(instr.value.property.identifier)) {
-              for (const operand of eachInstructionValueOperand(instr.value)) {
-                const hookKind = getHookKindForType(
-                  fn.env,
-                  instr.value.property.identifier.type,
-                );
-                if (hookKind != null) {
-                  validateNoRefValueAccess(errors, env, operand);
-                } else {
-                  validateNoRefAccess(errors, env, operand, operand.loc);
-                }
-              }
-            }
-            validateNoRefValueAccess(errors, env, instr.value.receiver);
-            const methType = env.get(instr.value.property.identifier.id);
-            let returnType: RefAccessType = {kind: 'None'};
-            if (methType?.kind === 'Structure' && methType.fn !== null) {
-              returnType = methType.fn.returnType;
-            }
-            env.set(instr.lvalue.identifier.id, returnType);
-            break;
-          }
+          case 'MethodCall':
           case 'CallExpression': {
-            const callee = instr.value.callee;
+            const callee =
+              instr.value.kind === 'CallExpression'
+                ? instr.value.callee
+                : instr.value.property;
             const hookKind = getHookKindForType(fn.env, callee.identifier.type);
-            const isUseEffect = isEffectHook(callee.identifier);
             let returnType: RefAccessType = {kind: 'None'};
-            if (!isUseEffect) {
-              // Report a more precise error when calling a local function that accesses a ref
-              const fnType = env.get(instr.value.callee.identifier.id);
-              if (fnType?.kind === 'Structure' && fnType.fn !== null) {
-                returnType = fnType.fn.returnType;
-                if (fnType.fn.readRefEffect) {
-                  errors.push({
-                    severity: ErrorSeverity.InvalidReact,
-                    reason:
-                      'This function accesses a ref value (the `current` property), which may not be accessed during render. (https://react.dev/reference/react/useRef)',
-                    loc: callee.loc,
-                    description:
-                      callee.identifier.name !== null &&
-                      callee.identifier.name.kind === 'named'
-                        ? `Function \`${callee.identifier.name.value}\` accesses a ref`
-                        : null,
-                    suggestions: null,
-                  });
-                }
+            const fnType = env.get(callee.identifier.id);
+            if (fnType?.kind === 'Structure' && fnType.fn !== null) {
+              returnType = fnType.fn.returnType;
+              if (fnType.fn.readRefEffect) {
+                errors.push({
+                  severity: ErrorSeverity.InvalidReact,
+                  reason:
+                    'This function accesses a ref value (the `current` property), which may not be accessed during render. (https://react.dev/reference/react/useRef)',
+                  loc: callee.loc,
+                  description:
+                    callee.identifier.name !== null &&
+                    callee.identifier.name.kind === 'named'
+                      ? `Function \`${callee.identifier.name.value}\` accesses a ref`
+                      : null,
+                  suggestions: null,
+                });
               }
-              for (const operand of eachInstructionValueOperand(instr.value)) {
-                if (hookKind != null) {
-                  validateNoRefValueAccess(errors, env, operand);
-                } else {
-                  validateNoRefAccess(errors, env, operand, operand.loc);
-                }
+            }
+            for (const operand of eachInstructionValueOperand(instr.value)) {
+              if (hookKind != null) {
+                validateNoDirectRefValueAccess(errors, operand, env);
+              } else {
+                validateNoRefAccess(errors, env, operand, operand.loc);
               }
             }
             env.set(instr.lvalue.identifier.id, returnType);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.hook-ref-value.expect.md

@@ -0,0 +1,34 @@
+
+## Input
+
+```javascript
+import {useEffect} from 'react';
+
+function Component(props) {
+  const ref = useRef();
+  useEffect(() => {}, [ref.current]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+
+
+## Error
+
+```
+  3 | function Component(props) {
+  4 |   const ref = useRef();
+> 5 |   useEffect(() => {}, [ref.current]);
+    |                        ^^^^^^^^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (5:5)
+
+InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (5:5)
+  6 | }
+  7 |
+  8 | export const FIXTURE_ENTRYPOINT = {
+```
+
+

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.hook-ref-value.js

@@ -0,0 +1,11 @@
+import {useEffect} from 'react';
+
+function Component(props) {
+  const ref = useRef();
+  useEffect(() => {}, [ref.current]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};

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

@@ -0,0 +1,54 @@
+
+## Input
+
+```javascript
+import {useEffect} from 'react';
+
+function Component(props) {
+  const ref = useRef();
+  useFoo(() => {
+    ref.current = 42;
+  });
+}
+
+function useFoo(x) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useEffect } from "react";
+
+function Component(props) {
+  const $ = _c(1);
+  const ref = useRef();
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      ref.current = 42;
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  useFoo(t0);
+}
+
+function useFoo(x) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+
+### Eval output
+(kind: exception) useRef is not defined

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hook-ref-callback.js

@@ -0,0 +1,15 @@
+import {useEffect} from 'react';
+
+function Component(props) {
+  const ref = useRef();
+  useFoo(() => {
+    ref.current = 42;
+  });
+}
+
+function useFoo(x) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useImperativeHandle-ref-mutate.expect.md

@@ -0,0 +1,66 @@
+
+## Input
+
+```javascript
+// @flow
+
+import {useImperativeHandle, useRef} from 'react';
+
+component Component(prop: number) {
+  const ref1 = useRef(null);
+  const ref2 = useRef(1);
+  useImperativeHandle(ref1, () => {
+    const precomputed = prop + ref2.current;
+    return {
+      foo: () => prop + ref2.current + precomputed,
+    };
+  }, [prop]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{prop: 1}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+
+import { useImperativeHandle, useRef } from "react";
+
+function Component(t0) {
+  const $ = _c(3);
+  const { prop } = t0;
+  const ref1 = useRef(null);
+  const ref2 = useRef(1);
+  let t1;
+  let t2;
+  if ($[0] !== prop) {
+    t1 = () => {
+      const precomputed = prop + ref2.current;
+      return { foo: () => prop + ref2.current + precomputed };
+    };
+
+    t2 = [prop];
+    $[0] = prop;
+    $[1] = t1;
+    $[2] = t2;
+  } else {
+    t1 = $[1];
+    t2 = $[2];
+  }
+  useImperativeHandle(ref1, t1, t2);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ prop: 1 }],
+};
+
+```
+
+### Eval output
+(kind: ok) 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useImperativeHandle-ref-mutate.js

@@ -0,0 +1,19 @@
+// @flow
+
+import {useImperativeHandle, useRef} from 'react';
+
+component Component(prop: number) {
+  const ref1 = useRef(null);
+  const ref2 = useRef(1);
+  useImperativeHandle(ref1, () => {
+    const precomputed = prop + ref2.current;
+    return {
+      foo: () => prop + ref2.current + precomputed,
+    };
+  }, [prop]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{prop: 1}],
+};