[compiler] Allow ref access in useImperativeHandle

Summary:
Code within the useImperativeHandle callback is not executed during render, much like useEffect. We can apply the same rules around ref access to this callback.

ghstack-source-id: 8e52839
Pull Request resolved: #30913

Author: mvitousek

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

@@ -282,7 +282,12 @@ function validateNoRefAccessInRenderImpl(
             break;
           }
           case 'MethodCall': {
-            if (!isEffectHook(instr.value.property.identifier)) {
+            if (
+              !isHookThatAllowsRefValueAccess(
+                fn,
+                instr.value.property.identifier,
+              )
+            ) {
               for (const operand of eachInstructionValueOperand(instr.value)) {
                 const hookKind = getHookKindForType(
                   fn.env,
@@ -307,7 +312,10 @@ function validateNoRefAccessInRenderImpl(
           case 'CallExpression': {
             const callee = instr.value.callee;
             const hookKind = getHookKindForType(fn.env, callee.identifier.type);
-            const isUseEffect = isEffectHook(callee.identifier);
+            const isUseEffect = isHookThatAllowsRefValueAccess(
+              fn,
+              callee.identifier,
+            );
             let returnType: RefAccessType = {kind: 'None'};
             if (!isUseEffect) {
               // Report a more precise error when calling a local function that accesses a ref
@@ -499,3 +507,11 @@ function validateNoDirectRefValueAccess(
     });
   }
 }
+
+function isHookThatAllowsRefValueAccess(
+  fn: HIRFunction,
+  identifier: Identifier,
+): boolean {
+  const hookKind = getHookKindForType(fn.env, identifier.type);
+  return isEffectHook(identifier) || hookKind === 'useImperativeHandle';
+}

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}],
+};