Update on "[compiler] InferReferenceEffects outputs a disjoint set of aliases"

Test Plan:
For the purposes of a pass added later in this stack, the InferReferenceEffects now outputs a DisjointSet of identifier ids representing aliases. We add this to InferReferenceEffects because this pass already does its own alias analysis, so making it available to later passes is convenient.

This entails implementing copy() and equals() methods on DIsjointSets so that the InferenceState copy and merge methods can handle alias sets.

[ghstack-poisoned]

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

@@ -364,6 +364,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/Inference/InferReferenceEffects.ts

@@ -227,6 +227,7 @@ export default function inferReferenceEffects(
   queue(fn.body.entry, initialState);
 
   const finishedStates: Map<BlockId, InferenceState> = new Map();
+  const functionEffects: Array<FunctionEffect> = fn.effects ?? [];
 
   while (queuedStates.size !== 0) {
     for (const [blockId, block] of fn.body.blocks) {
@@ -247,16 +248,17 @@ export default function inferReferenceEffects(
     }
   }
 
-  const summaryState = Array(...finishedStates.values()).reduce(
-    (acc, state) => acc.merge(state) ?? acc,
-  );
-
-  return summaryState.aliases;
   if (options.isFunctionExpression) {
     fn.effects = functionEffects;
   } else {
     raiseFunctionEffectErrors(functionEffects);
   }
+
+  const summaryState = Array(...finishedStates.values()).reduce(
+    (acc, state) => acc.merge(state) ?? acc,
+  );
+
+  return summaryState.aliases;
 }
 
 type FreezeAction = {values: Set<InstructionValue>; reason: Set<ValueReason>};

compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts

@@ -107,7 +107,7 @@ function equation(left: Type, right: Type): TypeEquation {
 function* generate(
   func: HIRFunction,
 ): Generator<TypeEquation, void, undefined> {
-  if (func.env.fnType === 'Component') {
+  if (func.fnType === 'Component') {
     const [props, ref] = func.params;
     if (props && props.kind === 'Identifier') {
       yield equation(props.identifier.type, {

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

@@ -8,9 +8,11 @@
 import {CompilerError, ErrorSeverity} from '../CompilerError';
 import {
   HIRFunction,
+  Identifier,
   IdentifierId,
   Place,
   SourceLocation,
+  getHookKindForType,
   isRefValueType,
   isUseRefType,
 } from '../HIR';
@@ -20,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
@@ -42,25 +43,154 @@ import {isEffectHook} from './ValidateMemoizedEffectDependencies';
  * In the future we may reject more cases, based on either object names (`fooRef.current` is likely a ref)
  * or based on property name alone (`foo.current` might be a ref).
  */
-type State = {
-  refs: Set<IdentifierId>;
-  refValues: Map<IdentifierId, SourceLocation | null>;
-  refAccessingFunctions: Set<IdentifierId>;
-};
+
+type RefAccessType = {kind: 'None'} | RefAccessRefType;
+
+type RefAccessRefType =
+  | {kind: 'Ref'}
+  | {kind: 'RefValue'; loc?: SourceLocation}
+  | {kind: 'Structure'; value: null | RefAccessRefType; fn: null | RefFnType};
+
+type RefFnType = {readRefEffect: boolean; returnType: RefAccessType};
+
+class Env extends Map<IdentifierId, RefAccessType> {
+  #changed = false;
+
+  resetChanged(): void {
+    this.#changed = false;
+  }
+
+  hasChanged(): boolean {
+    return this.#changed;
+  }
+
+  override set(key: IdentifierId, value: RefAccessType): this {
+    const cur = this.get(key);
+    const widenedValue = joinRefAccessTypes(value, cur ?? {kind: 'None'});
+    if (
+      !(cur == null && widenedValue.kind === 'None') &&
+      (cur == null || !tyEqual(cur, widenedValue))
+    ) {
+      this.#changed = true;
+    }
+    return super.set(key, widenedValue);
+  }
+}
 
 export function validateNoRefAccessInRender(fn: HIRFunction): void {
-  const state = {
-    refs: new Set<IdentifierId>(),
-    refValues: new Map<IdentifierId, SourceLocation | null>(),
-    refAccessingFunctions: new Set<IdentifierId>(),
-  };
-  validateNoRefAccessInRenderImpl(fn, state).unwrap();
+  const env = new Env();
+  validateNoRefAccessInRenderImpl(fn, env).unwrap();
+}
+
+function refTypeOfType(identifier: Identifier): RefAccessType {
+  if (isRefValueType(identifier)) {
+    return {kind: 'RefValue'};
+  } else if (isUseRefType(identifier)) {
+    return {kind: 'Ref'};
+  } else {
+    return {kind: 'None'};
+  }
+}
+
+function tyEqual(a: RefAccessType, b: RefAccessType): boolean {
+  if (a.kind !== b.kind) {
+    return false;
+  }
+  switch (a.kind) {
+    case 'None':
+      return true;
+    case 'Ref':
+      return true;
+    case 'RefValue':
+      CompilerError.invariant(b.kind === 'RefValue', {
+        reason: 'Expected ref value',
+        loc: null,
+      });
+      return a.loc == b.loc;
+    case 'Structure': {
+      CompilerError.invariant(b.kind === 'Structure', {
+        reason: 'Expected structure',
+        loc: null,
+      });
+      const fnTypesEqual =
+        (a.fn === null && b.fn === null) ||
+        (a.fn !== null &&
+          b.fn !== null &&
+          a.fn.readRefEffect === b.fn.readRefEffect &&
+          tyEqual(a.fn.returnType, b.fn.returnType));
+      return (
+        fnTypesEqual &&
+        (a.value === b.value ||
+          (a.value !== null && b.value !== null && tyEqual(a.value, b.value)))
+      );
+    }
+  }
+}
+
+function joinRefAccessTypes(...types: Array<RefAccessType>): RefAccessType {
+  function joinRefAccessRefTypes(
+    a: RefAccessRefType,
+    b: RefAccessRefType,
+  ): RefAccessRefType {
+    if (a.kind === 'RefValue') {
+      return a;
+    } else if (b.kind === 'RefValue') {
+      return b;
+    } else if (a.kind === 'Ref' || b.kind === 'Ref') {
+      return {kind: 'Ref'};
+    } else {
+      CompilerError.invariant(
+        a.kind === 'Structure' && b.kind === 'Structure',
+        {
+          reason: 'Expected structure',
+          loc: null,
+        },
+      );
+      const fn =
+        a.fn === null
+          ? b.fn
+          : b.fn === null
+            ? a.fn
+            : {
+                readRefEffect: a.fn.readRefEffect || b.fn.readRefEffect,
+                returnType: joinRefAccessTypes(
+                  a.fn.returnType,
+                  b.fn.returnType,
+                ),
+              };
+      const value =
+        a.value === null
+          ? b.value
+          : b.value === null
+            ? a.value
+            : joinRefAccessRefTypes(a.value, b.value);
+      return {
+        kind: 'Structure',
+        fn,
+        value,
+      };
+    }
+  }
+
+  return types.reduce(
+    (a, b) => {
+      if (a.kind === 'None') {
+        return b;
+      } else if (b.kind === 'None') {
+        return a;
+      } else {
+        return joinRefAccessRefTypes(a, b);
+      }
+    },
+    {kind: 'None'},
+  );
 }
 
 function validateNoRefAccessInRenderImpl(
   fn: HIRFunction,
-  state: State,
-): Result<void, CompilerError> {
+  env: Env,
+): Result<RefAccessType, CompilerError> {
+  let returnValues: Array<undefined | RefAccessType> = [];
   let place;
   for (const param of fn.params) {
     if (param.kind === 'Identifier') {
@@ -68,293 +198,266 @@ function validateNoRefAccessInRenderImpl(
     } else {
       place = param.place;
     }
-
-    if (isRefValueType(place.identifier)) {
-      state.refValues.set(place.identifier.id, null);
-    }
-    if (isUseRefType(place.identifier)) {
-      state.refs.add(place.identifier.id);
-    }
+    const type = refTypeOfType(place.identifier);
+    env.set(place.identifier.id, type);
   }
-  const errors = new CompilerError();
-  for (const [, block] of fn.body.blocks) {
-    for (const phi of block.phis) {
-      phi.operands.forEach(operand => {
-        if (state.refs.has(operand.id) || isUseRefType(phi.id)) {
-          state.refs.add(phi.id.id);
-        }
-        const refValue = state.refValues.get(operand.id);
-        if (refValue !== undefined || isRefValueType(operand)) {
-          state.refValues.set(
-            phi.id.id,
-            refValue ?? state.refValues.get(phi.id.id) ?? null,
-          );
-        }
-        if (state.refAccessingFunctions.has(operand.id)) {
-          state.refAccessingFunctions.add(phi.id.id);
-        }
-      });
-    }
 
-    for (const instr of block.instructions) {
-      for (const operand of eachInstructionValueOperand(instr.value)) {
-        if (isRefValueType(operand.identifier)) {
-          CompilerError.invariant(state.refValues.has(operand.identifier.id), {
-            reason: 'Expected ref value to be in state',
-            loc: operand.loc,
-          });
-        }
-        if (isUseRefType(operand.identifier)) {
-          CompilerError.invariant(state.refs.has(operand.identifier.id), {
-            reason: 'Expected ref to be in state',
-            loc: operand.loc,
-          });
-        }
+  for (let i = 0; (i == 0 || env.hasChanged()) && i < 10; i++) {
+    env.resetChanged();
+    returnValues = [];
+    const errors = new CompilerError();
+    for (const [, block] of fn.body.blocks) {
+      for (const phi of block.phis) {
+        env.set(
+          phi.id.id,
+          joinRefAccessTypes(
+            ...Array(...phi.operands.values()).map(
+              operand => env.get(operand.id) ?? ({kind: 'None'} as const),
+            ),
+          ),
+        );
       }
 
-      switch (instr.value.kind) {
-        case 'JsxExpression':
-        case 'JsxFragment': {
-          for (const operand of eachInstructionValueOperand(instr.value)) {
-            validateNoDirectRefValueAccess(errors, operand, state);
-          }
-          break;
-        }
-        case 'ComputedLoad':
-        case 'PropertyLoad': {
-          if (typeof instr.value.property !== 'string') {
-            validateNoRefValueAccess(errors, state, instr.value.property);
-          }
-          if (
-            state.refAccessingFunctions.has(instr.value.object.identifier.id)
-          ) {
-            state.refAccessingFunctions.add(instr.lvalue.identifier.id);
-          }
-          if (state.refs.has(instr.value.object.identifier.id)) {
-            /*
-             * Once an object contains a ref at any level, we treat it as a ref.
-             * If we look something up from it, that value may either be a ref
-             * or the ref value (or neither), so we conservatively assume it's both.
-             */
-            state.refs.add(instr.lvalue.identifier.id);
-            state.refValues.set(instr.lvalue.identifier.id, instr.loc);
-          }
-          break;
-        }
-        case 'LoadContext':
-        case 'LoadLocal': {
-          if (
-            state.refAccessingFunctions.has(instr.value.place.identifier.id)
-          ) {
-            state.refAccessingFunctions.add(instr.lvalue.identifier.id);
-          }
-          const refValue = state.refValues.get(instr.value.place.identifier.id);
-          if (refValue !== undefined) {
-            state.refValues.set(instr.lvalue.identifier.id, refValue);
-          }
-          if (state.refs.has(instr.value.place.identifier.id)) {
-            state.refs.add(instr.lvalue.identifier.id);
-          }
-          break;
-        }
-        case 'StoreContext':
-        case 'StoreLocal': {
-          if (
-            state.refAccessingFunctions.has(instr.value.value.identifier.id)
-          ) {
-            state.refAccessingFunctions.add(
-              instr.value.lvalue.place.identifier.id,
-            );
-            state.refAccessingFunctions.add(instr.lvalue.identifier.id);
-          }
-          const refValue = state.refValues.get(instr.value.value.identifier.id);
-          if (
-            refValue !== undefined ||
-            isRefValueType(instr.value.lvalue.place.identifier)
-          ) {
-            state.refValues.set(
-              instr.value.lvalue.place.identifier.id,
-              refValue ?? null,
-            );
-            state.refValues.set(instr.lvalue.identifier.id, refValue ?? null);
-          }
-          if (state.refs.has(instr.value.value.identifier.id)) {
-            state.refs.add(instr.value.lvalue.place.identifier.id);
-            state.refs.add(instr.lvalue.identifier.id);
-          }
-          break;
-        }
-        case 'Destructure': {
-          const destructuredFunction = state.refAccessingFunctions.has(
-            instr.value.value.identifier.id,
-          );
-          const destructuredRef = state.refs.has(
-            instr.value.value.identifier.id,
-          );
-          for (const lval of eachPatternOperand(instr.value.lvalue.pattern)) {
-            if (isUseRefType(lval.identifier)) {
-              state.refs.add(lval.identifier.id);
-            }
-            if (destructuredRef || isRefValueType(lval.identifier)) {
-              state.refs.add(lval.identifier.id);
-              state.refValues.set(lval.identifier.id, null);
-            }
-            if (destructuredFunction) {
-              state.refAccessingFunctions.add(lval.identifier.id);
-            }
-          }
-          break;
-        }
-        case 'ObjectMethod':
-        case 'FunctionExpression': {
-          if (
-            /*
-             * check if the function expression accesses a ref *or* some other
-             * function which accesses a ref
-             */
-            [...eachInstructionValueOperand(instr.value)].some(
-              operand =>
-                state.refValues.has(operand.identifier.id) ||
-                state.refAccessingFunctions.has(operand.identifier.id),
-            ) ||
-            // check for cases where .current is accessed through an aliased ref
-            ([...eachInstructionValueOperand(instr.value)].some(operand =>
-              state.refs.has(operand.identifier.id),
-            ) &&
-              validateNoRefAccessInRenderImpl(
-                instr.value.loweredFunc.func,
-                state,
-              ).isErr())
-          ) {
-            // This function expression unconditionally accesses a ref
-            state.refAccessingFunctions.add(instr.lvalue.identifier.id);
-          }
-          break;
-        }
-        case 'MethodCall': {
-          if (!isEffectHook(instr.value.property.identifier)) {
+      for (const instr of block.instructions) {
+        switch (instr.value.kind) {
+          case 'JsxExpression':
+          case 'JsxFragment': {
             for (const operand of eachInstructionValueOperand(instr.value)) {
-              validateNoRefAccess(errors, state, operand, operand.loc);
+              validateNoDirectRefValueAccess(errors, operand, env);
             }
+            break;
           }
-          break;
-        }
-        case 'CallExpression': {
-          const callee = instr.value.callee;
-          const isUseEffect = isEffectHook(callee.identifier);
-          if (!isUseEffect) {
-            // Report a more precise error when calling a local function that accesses a ref
-            if (state.refAccessingFunctions.has(callee.identifier.id)) {
-              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,
-              });
+          case 'ComputedLoad':
+          case 'PropertyLoad': {
+            if (typeof instr.value.property !== 'string') {
+              validateNoDirectRefValueAccess(errors, instr.value.property, env);
             }
-            for (const operand of eachInstructionValueOperand(instr.value)) {
-              validateNoRefAccess(
-                errors,
-                state,
-                operand,
-                state.refValues.get(operand.identifier.id) ?? operand.loc,
+            const objType = env.get(instr.value.object.identifier.id);
+            let lookupType: null | RefAccessType = null;
+            if (objType?.kind === 'Structure') {
+              lookupType = objType.value;
+            } else if (objType?.kind === 'Ref') {
+              lookupType = {kind: 'RefValue', loc: instr.loc};
+            }
+            env.set(
+              instr.lvalue.identifier.id,
+              lookupType ?? refTypeOfType(instr.lvalue.identifier),
+            );
+            break;
+          }
+          case 'LoadContext':
+          case 'LoadLocal': {
+            env.set(
+              instr.lvalue.identifier.id,
+              env.get(instr.value.place.identifier.id) ??
+                refTypeOfType(instr.lvalue.identifier),
+            );
+            break;
+          }
+          case 'StoreContext':
+          case 'StoreLocal': {
+            env.set(
+              instr.value.lvalue.place.identifier.id,
+              env.get(instr.value.value.identifier.id) ??
+                refTypeOfType(instr.value.lvalue.place.identifier),
+            );
+            env.set(
+              instr.lvalue.identifier.id,
+              env.get(instr.value.value.identifier.id) ??
+                refTypeOfType(instr.lvalue.identifier),
+            );
+            break;
+          }
+          case 'Destructure': {
+            const objType = env.get(instr.value.value.identifier.id);
+            let lookupType = null;
+            if (objType?.kind === 'Structure') {
+              lookupType = objType.value;
+            }
+            env.set(
+              instr.lvalue.identifier.id,
+              lookupType ?? refTypeOfType(instr.lvalue.identifier),
+            );
+            for (const lval of eachPatternOperand(instr.value.lvalue.pattern)) {
+              env.set(
+                lval.identifier.id,
+                lookupType ?? refTypeOfType(lval.identifier),
               );
             }
+            break;
           }
-          break;
-        }
-        case 'ObjectExpression':
-        case 'ArrayExpression': {
-          for (const operand of eachInstructionValueOperand(instr.value)) {
-            validateNoDirectRefValueAccess(errors, operand, state);
-            if (state.refAccessingFunctions.has(operand.identifier.id)) {
-              state.refAccessingFunctions.add(instr.lvalue.identifier.id);
-            }
-            if (state.refs.has(operand.identifier.id)) {
-              state.refs.add(instr.lvalue.identifier.id);
-            }
-            const refValue = state.refValues.get(operand.identifier.id);
-            if (refValue !== undefined) {
-              state.refValues.set(instr.lvalue.identifier.id, refValue);
+          case 'ObjectMethod':
+          case 'FunctionExpression': {
+            let returnType: RefAccessType = {kind: 'None'};
+            let readRefEffect = false;
+            const result = validateNoRefAccessInRenderImpl(
+              instr.value.loweredFunc.func,
+              env,
+            );
+            if (result.isOk()) {
+              returnType = result.unwrap();
+            } else if (result.isErr()) {
+              readRefEffect = true;
             }
+            env.set(instr.lvalue.identifier.id, {
+              kind: 'Structure',
+              fn: {
+                readRefEffect,
+                returnType,
+              },
+              value: null,
+            });
+            break;
+          }
+          case 'MethodCall':
+          case 'CallExpression': {
+            const callee =
+              instr.value.kind === 'CallExpression'
+                ? instr.value.callee
+                : instr.value.property;
+            const hookKind = getHookKindForType(fn.env, callee.identifier.type);
+            let returnType: RefAccessType = {kind: 'None'};
+            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) {
+                validateNoDirectRefValueAccess(errors, operand, env);
+              } else {
+                validateNoRefAccess(errors, env, operand, operand.loc);
+              }
+            }
+            env.set(instr.lvalue.identifier.id, returnType);
+            break;
+          }
+          case 'ObjectExpression':
+          case 'ArrayExpression': {
+            const types: Array<RefAccessType> = [];
+            for (const operand of eachInstructionValueOperand(instr.value)) {
+              validateNoDirectRefValueAccess(errors, operand, env);
+              types.push(env.get(operand.identifier.id) ?? {kind: 'None'});
+            }
+            const value = joinRefAccessTypes(...types);
+            if (value.kind === 'None') {
+              env.set(instr.lvalue.identifier.id, {kind: 'None'});
+            } else {
+              env.set(instr.lvalue.identifier.id, {
+                kind: 'Structure',
+                value,
+                fn: null,
+              });
+            }
+            break;
+          }
+          case 'PropertyDelete':
+          case 'PropertyStore':
+          case 'ComputedDelete':
+          case 'ComputedStore': {
+            validateNoRefAccess(errors, env, instr.value.object, instr.loc);
+            for (const operand of eachInstructionValueOperand(instr.value)) {
+              if (operand === instr.value.object) {
+                continue;
+              }
+              validateNoRefValueAccess(errors, env, operand);
+            }
+            break;
+          }
+          case 'StartMemoize':
+          case 'FinishMemoize':
+            break;
+          default: {
+            for (const operand of eachInstructionValueOperand(instr.value)) {
+              validateNoRefValueAccess(errors, env, operand);
+            }
+            break;
           }
-          break;
         }
-        case 'PropertyDelete':
-        case 'PropertyStore':
-        case 'ComputedDelete':
-        case 'ComputedStore': {
-          validateNoRefAccess(
-            errors,
-            state,
-            instr.value.object,
-            state.refValues.get(instr.value.object.identifier.id) ?? instr.loc,
+        if (isUseRefType(instr.lvalue.identifier)) {
+          env.set(
+            instr.lvalue.identifier.id,
+            joinRefAccessTypes(
+              env.get(instr.lvalue.identifier.id) ?? {kind: 'None'},
+              {kind: 'Ref'},
+            ),
           );
-          for (const operand of eachInstructionValueOperand(instr.value)) {
-            if (operand === instr.value.object) {
-              continue;
-            }
-            validateNoRefValueAccess(errors, state, operand);
-          }
-          break;
         }
-        case 'StartMemoize':
-        case 'FinishMemoize':
-          break;
-        default: {
-          for (const operand of eachInstructionValueOperand(instr.value)) {
-            validateNoRefValueAccess(errors, state, operand);
-          }
-          break;
+        if (isRefValueType(instr.lvalue.identifier)) {
+          env.set(
+            instr.lvalue.identifier.id,
+            joinRefAccessTypes(
+              env.get(instr.lvalue.identifier.id) ?? {kind: 'None'},
+              {kind: 'RefValue', loc: instr.loc},
+            ),
+          );
         }
       }
-      if (isUseRefType(instr.lvalue.identifier)) {
-        state.refs.add(instr.lvalue.identifier.id);
-      }
-      if (
-        isRefValueType(instr.lvalue.identifier) &&
-        !state.refValues.has(instr.lvalue.identifier.id)
-      ) {
-        state.refValues.set(instr.lvalue.identifier.id, instr.loc);
+      for (const operand of eachTerminalOperand(block.terminal)) {
+        if (block.terminal.kind !== 'return') {
+          validateNoRefValueAccess(errors, env, operand);
+        } else {
+          // Allow functions containing refs to be returned, but not direct ref values
+          validateNoDirectRefValueAccess(errors, operand, env);
+          returnValues.push(env.get(operand.identifier.id));
+        }
       }
     }
-    for (const operand of eachTerminalOperand(block.terminal)) {
-      if (block.terminal.kind !== 'return') {
-        validateNoRefValueAccess(errors, state, operand);
-      } else {
-        // Allow functions containing refs to be returned, but not direct ref values
-        validateNoDirectRefValueAccess(errors, operand, state);
-      }
+
+    if (errors.hasErrors()) {
+      return Err(errors);
     }
   }
 
-  if (errors.hasErrors()) {
-    return Err(errors);
-  } else {
-    return Ok(undefined);
+  CompilerError.invariant(!env.hasChanged(), {
+    reason: 'Ref type environment did not converge',
+    loc: null,
+  });
+
+  return Ok(
+    joinRefAccessTypes(
+      ...returnValues.filter((env): env is RefAccessType => env !== undefined),
+    ),
+  );
+}
+
+function destructure(
+  type: RefAccessType | undefined,
+): RefAccessType | undefined {
+  if (type?.kind === 'Structure' && type.value !== null) {
+    return destructure(type.value);
   }
+  return type;
 }
 
 function validateNoRefValueAccess(
   errors: CompilerError,
-  state: State,
+  env: Env,
   operand: Place,
 ): void {
+  const type = destructure(env.get(operand.identifier.id));
   if (
-    state.refValues.has(operand.identifier.id) ||
-    state.refAccessingFunctions.has(operand.identifier.id)
+    type?.kind === 'RefValue' ||
+    (type?.kind === 'Structure' && type.fn?.readRefEffect)
   ) {
     errors.push({
       severity: ErrorSeverity.InvalidReact,
       reason:
         'Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef)',
-      loc: state.refValues.get(operand.identifier.id) ?? operand.loc,
+      loc: (type.kind === 'RefValue' && type.loc) || operand.loc,
       description:
         operand.identifier.name !== null &&
         operand.identifier.name.kind === 'named'
@@ -367,20 +470,21 @@ function validateNoRefValueAccess(
 
 function validateNoRefAccess(
   errors: CompilerError,
-  state: State,
+  env: Env,
   operand: Place,
   loc: SourceLocation,
 ): void {
+  const type = destructure(env.get(operand.identifier.id));
   if (
-    state.refs.has(operand.identifier.id) ||
-    state.refValues.has(operand.identifier.id) ||
-    state.refAccessingFunctions.has(operand.identifier.id)
+    type?.kind === 'Ref' ||
+    type?.kind === 'RefValue' ||
+    (type?.kind === 'Structure' && type.fn?.readRefEffect)
   ) {
     errors.push({
       severity: ErrorSeverity.InvalidReact,
       reason:
         'Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef)',
-      loc: loc,
+      loc: (type.kind === 'RefValue' && type.loc) || loc,
       description:
         operand.identifier.name !== null &&
         operand.identifier.name.kind === 'named'
@@ -394,14 +498,15 @@ function validateNoRefAccess(
 function validateNoDirectRefValueAccess(
   errors: CompilerError,
   operand: Place,
-  state: State,
+  env: Env,
 ): void {
-  if (state.refValues.has(operand.identifier.id)) {
+  const type = destructure(env.get(operand.identifier.id));
+  if (type?.kind === 'RefValue') {
     errors.push({
       severity: ErrorSeverity.InvalidReact,
       reason:
         'Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef)',
-      loc: state.refValues.get(operand.identifier.id) ?? operand.loc,
+      loc: type.loc ?? operand.loc,
       description:
         operand.identifier.name !== null &&
         operand.identifier.name.kind === 'named'

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/capture-ref-for-later-mutation.expect.md

@@ -0,0 +1,69 @@
+
+## Input
+
+```javascript
+import {useRef} from 'react';
+import {addOne} from 'shared-runtime';
+
+function useKeyCommand() {
+  const currentPosition = useRef(0);
+  const handleKey = direction => () => {
+    const position = currentPosition.current;
+    const nextPosition = direction === 'left' ? addOne(position) : position;
+    currentPosition.current = nextPosition;
+  };
+  const moveLeft = {
+    handler: handleKey('left'),
+  };
+  const moveRight = {
+    handler: handleKey('right'),
+  };
+  return [moveLeft, moveRight];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useKeyCommand,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useRef } from "react";
+import { addOne } from "shared-runtime";
+
+function useKeyCommand() {
+  const $ = _c(1);
+  const currentPosition = useRef(0);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const handleKey = (direction) => () => {
+      const position = currentPosition.current;
+      const nextPosition = direction === "left" ? addOne(position) : position;
+      currentPosition.current = nextPosition;
+    };
+
+    const moveLeft = { handler: handleKey("left") };
+
+    const moveRight = { handler: handleKey("right") };
+
+    t0 = [moveLeft, moveRight];
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useKeyCommand,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) [{"handler":"[[ function params=0 ]]"},{"handler":"[[ function params=0 ]]"}]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.capture-ref-for-mutation.expect.md

@@ -13,10 +13,10 @@ function useKeyCommand() {
     currentPosition.current = nextPosition;
   };
   const moveLeft = {
-    handler: handleKey('left'),
+    handler: handleKey('left')(),
   };
   const moveRight = {
-    handler: handleKey('right'),
+    handler: handleKey('right')(),
   };
   return [moveLeft, moveRight];
 }
@@ -34,8 +34,8 @@ export const FIXTURE_ENTRYPOINT = {
 ```
   10 |   };
   11 |   const moveLeft = {
-> 12 |     handler: handleKey('left'),
-     |              ^^^^^^^^^ InvalidReact: This function accesses a ref value (the `current` property), which may not be accessed during render. (https://react.dev/reference/react/useRef) (12:12)
+> 12 |     handler: handleKey('left')(),
+     |              ^^^^^^^^^^^^^^^^^ InvalidReact: This function accesses a ref value (the `current` property), which may not be accessed during render. (https://react.dev/reference/react/useRef) (12:12)
 
 InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (12:12)
 
@@ -44,7 +44,7 @@ InvalidReact: This function accesses a ref value (the `current` property), which
 InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (15:15)
   13 |   };
   14 |   const moveRight = {
-  15 |     handler: handleKey('right'),
+  15 |     handler: handleKey('right')(),
 ```
           
       

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.capture-ref-for-mutation.tsx

@@ -0,0 +1,23 @@
+import {useRef} from 'react';
+import {addOne} from 'shared-runtime';
+
+function useKeyCommand() {
+  const currentPosition = useRef(0);
+  const handleKey = direction => () => {
+    const position = currentPosition.current;
+    const nextPosition = direction === 'left' ? addOne(position) : position;
+    currentPosition.current = nextPosition;
+  };
+  const moveLeft = {
+    handler: handleKey('left')(),
+  };
+  const moveRight = {
+    handler: handleKey('right')(),
+  };
+  return [moveLeft, moveRight];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useKeyCommand,
+  params: [],
+};

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

@@ -0,0 +1,34 @@
+
+## Input
+
+```javascript
+import {useEffect, useRef} 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, useRef} 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, useRef} 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, useRef } 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: ok) 

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

@@ -0,0 +1,15 @@
+import {useEffect, useRef} 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/original-reactive-scopes-fork/capture-ref-for-later-mutation.expect.md

@@ -0,0 +1,70 @@
+
+## Input
+
+```javascript
+// @enableReactiveScopesInHIR:false
+import {useRef} from 'react';
+import {addOne} from 'shared-runtime';
+
+function useKeyCommand() {
+  const currentPosition = useRef(0);
+  const handleKey = direction => () => {
+    const position = currentPosition.current;
+    const nextPosition = direction === 'left' ? addOne(position) : position;
+    currentPosition.current = nextPosition;
+  };
+  const moveLeft = {
+    handler: handleKey('left'),
+  };
+  const moveRight = {
+    handler: handleKey('right'),
+  };
+  return [moveLeft, moveRight];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useKeyCommand,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableReactiveScopesInHIR:false
+import { useRef } from "react";
+import { addOne } from "shared-runtime";
+
+function useKeyCommand() {
+  const $ = _c(1);
+  const currentPosition = useRef(0);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const handleKey = (direction) => () => {
+      const position = currentPosition.current;
+      const nextPosition = direction === "left" ? addOne(position) : position;
+      currentPosition.current = nextPosition;
+    };
+
+    const moveLeft = { handler: handleKey("left") };
+
+    const moveRight = { handler: handleKey("right") };
+
+    t0 = [moveLeft, moveRight];
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useKeyCommand,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) [{"handler":"[[ function params=0 ]]"},{"handler":"[[ function params=0 ]]"}]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/original-reactive-scopes-fork/capture-ref-for-later-mutation.tsx

@@ -0,0 +1,24 @@
+// @enableReactiveScopesInHIR:false
+import {useRef} from 'react';
+import {addOne} from 'shared-runtime';
+
+function useKeyCommand() {
+  const currentPosition = useRef(0);
+  const handleKey = direction => () => {
+    const position = currentPosition.current;
+    const nextPosition = direction === 'left' ? addOne(position) : position;
+    currentPosition.current = nextPosition;
+  };
+  const moveLeft = {
+    handler: handleKey('left'),
+  };
+  const moveRight = {
+    handler: handleKey('right'),
+  };
+  return [moveLeft, moveRight];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useKeyCommand,
+  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}],
+};