[compiler] Alternate take on ref validation

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -85,7 +85,7 @@
   validateHooksUsage,
   validateMemoizedEffectDependencies,
   validateNoCapitalizedCalls,
   validateNoRefAccessInRender,
   validateNoSetStateInRender,
   validatePreservedManualMemoization,
   validateUseMemo,
@@ -285,7 +285,7 @@
     }
 
     if (env.config.validateRefAccessDuringRender) {
-      validateNoRefAccessInRender(hir).unwrap();
+      // validateNoRefAccessInRender(hir).unwrap();
     }
 
     if (env.config.validateNoSetStateInRender) {

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

@@ -221,7 +221,6 @@ export function lower(
     params,
     fnType: bindings == null ? env.fnType : 'Other',
     returnTypeAnnotation: null, // TODO: extract the actual return type node if present
-    returnType: makeType(),
     returns: createTemporaryPlace(env, func.node.loc ?? GeneratedSource),
     body: builder.build(),
     context,

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

@@ -279,7 +279,6 @@ export type HIRFunction = {
   env: Environment;
   params: Array<Place | SpreadPattern>;
   returnTypeAnnotation: t.FlowType | t.TSType | null;
-  returnType: Type;
   returns: Place;
   context: Array<Place>;
   effects: Array<FunctionEffect> | null;

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

@@ -54,6 +54,8 @@ export function printFunction(fn: HIRFunction): string {
   let definition = '';
   if (fn.id !== null) {
     definition += fn.id;
+  } else {
+    definition += '<<anonymous>>';
   }
   if (fn.params.length !== 0) {
     definition +=
@@ -71,10 +73,8 @@ export function printFunction(fn: HIRFunction): string {
   } else {
     definition += '()';
   }
-  if (definition.length !== 0) {
-    output.push(definition);
-  }
-  output.push(`: ${printType(fn.returnType)} @ ${printPlace(fn.returns)}`);
+  definition += `: ${printPlace(fn.returns)}`;
+  output.push(definition);
   output.push(...fn.directives);
   output.push(printHIR(fn.body));
   return output.join('\n');

compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts

@@ -28,7 +28,9 @@ import {
   isMapType,
   isPrimitiveType,
   isRefOrRefValue,
+  isRefValueType,
   isSetType,
+  isUseRefType,
   makeIdentifierId,
   Phi,
   Place,
@@ -219,6 +221,9 @@ export function inferMutationAliasingEffects(
       }
     }
   }
+  if (fn.env.config.validateRefAccessDuringRender) {
+    inferRefAccessEffects(fn, isFunctionExpression);
+  }
   return Ok(undefined);
 }
 
@@ -2513,3 +2518,127 @@ export type AbstractValue = {
   kind: ValueKind;
   reason: ReadonlySet<ValueReason>;
 };
+
+function inferRefAccessEffects(
+  fn: HIRFunction,
+  _isFunctionExpression: boolean,
+): void {
+  const nullish = new Set<IdentifierId>();
+  const nullishTest = new Map<IdentifierId, Place>();
+  let guard: {ref: IdentifierId; fallthrough: BlockId} | null = null;
+  const temporaries: Map<IdentifierId, Place> = new Map();
+
+  function visitOperand(operand: Place): AliasingEffect | null {
+    const nullTestRef = nullishTest.get(operand.identifier.id);
+    if (isRefValueType(operand.identifier) || nullTestRef != null) {
+      const refOperand = nullTestRef ?? operand;
+      return {
+        kind: 'Impure',
+        error: {
+          severity: ErrorSeverity.InvalidReact,
+          reason:
+            'Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef)',
+          loc: refOperand.loc,
+          description:
+            refOperand.identifier.name !== null &&
+            refOperand.identifier.name.kind === 'named'
+              ? `Cannot access ref value \`${refOperand.identifier.name.value}\``
+              : null,
+          suggestions: null,
+        },
+        place: refOperand,
+      };
+    }
+    return null;
+  }
+
+  for (const block of fn.body.blocks.values()) {
+    if (guard !== null && guard.fallthrough === block.id) {
+      guard = null;
+    }
+    for (const instr of block.instructions) {
+      const {lvalue, value} = instr;
+      if (value.kind === 'LoadLocal' && isUseRefType(value.place.identifier)) {
+        temporaries.set(lvalue.identifier.id, value.place);
+      } else if (
+        value.kind === 'StoreLocal' &&
+        isUseRefType(value.value.identifier)
+      ) {
+        temporaries.set(value.lvalue.place.identifier.id, value.value);
+        temporaries.set(lvalue.identifier.id, value.value);
+      } else if (
+        value.kind === 'BinaryExpression' &&
+        ((isRefValueType(value.left.identifier) &&
+          nullish.has(value.right.identifier.id)) ||
+          (nullish.has(value.left.identifier.id) &&
+            isRefValueType(value.right.identifier)))
+      ) {
+        const refOperand = isRefValueType(value.left.identifier)
+          ? value.left
+          : value.right;
+        const operand = temporaries.get(refOperand.identifier.id) ?? refOperand;
+        nullishTest.set(lvalue.identifier.id, operand);
+      } else if (value.kind === 'Primitive' && value.value == null) {
+        nullish.add(lvalue.identifier.id);
+      } else if (
+        value.kind === 'PropertyLoad' &&
+        isUseRefType(value.object.identifier) &&
+        value.property === 'current'
+      ) {
+        const refOperand =
+          temporaries.get(value.object.identifier.id) ?? value.object;
+        temporaries.set(lvalue.identifier.id, refOperand);
+      } else if (
+        value.kind === 'PropertyStore' &&
+        value.property === 'current' &&
+        isUseRefType(value.object.identifier)
+      ) {
+        const refOperand =
+          temporaries.get(value.object.identifier.id) ?? value.object;
+        if (guard != null && refOperand.identifier.id === guard.ref) {
+          // Allow a single write within the guard
+          guard = null;
+        } else {
+          instr.effects ??= [];
+          instr.effects.push({
+            kind: 'Impure',
+            error: {
+              severity: ErrorSeverity.InvalidReact,
+              reason:
+                'Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef)',
+              loc: value.loc,
+              description:
+                value.object.identifier.name !== null &&
+                value.object.identifier.name.kind === 'named'
+                  ? `Cannot access ref value \`${value.object.identifier.name.value}\``
+                  : null,
+              suggestions: null,
+            },
+            place: value.object,
+          });
+        }
+      } else {
+        for (const operand of eachInstructionValueOperand(value)) {
+          const error = visitOperand(operand);
+          if (error) {
+            instr.effects ??= [];
+            instr.effects.push(error);
+          }
+        }
+      }
+    }
+    if (
+      guard == null &&
+      block.terminal.kind === 'if' &&
+      nullishTest.has(block.terminal.test.identifier.id)
+    ) {
+      const ref = nullishTest.get(block.terminal.test.identifier.id)!;
+      guard = {ref: ref.identifier.id, fallthrough: block.terminal.fallthrough};
+    } else {
+      for (const operand of eachTerminalOperand(block.terminal)) {
+        const _effect = visitOperand(operand);
+        // TODO: need a place to store terminal effects generically
+      }
+    }
+  }
+}

compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts

@@ -18,6 +18,7 @@ import {
   ValueKind,
   ValueReason,
   Place,
+  isPrimitiveType,
 } from '../HIR/HIR';
 import {
   eachInstructionLValue,
@@ -471,15 +472,15 @@ export function inferMutationAliasingRanges(
    * Here we populate an effect to create the return value as well as populating alias/capture
    * effects for how data flows between the params, context vars, and return.
    */
+  const returns = fn.returns.identifier;
   functionEffects.push({
     kind: 'Create',
     into: fn.returns,
-    value:
-      fn.returnType.kind === 'Primitive'
-        ? ValueKind.Primitive
-        : isJsxType(fn.returnType)
-          ? ValueKind.Frozen
-          : ValueKind.Mutable,
+    value: isPrimitiveType(returns)
+      ? ValueKind.Primitive
+      : isJsxType(returns.type)
+        ? ValueKind.Frozen
+        : ValueKind.Mutable,
     reason: ValueReason.KnownReturnSignature,
   });
   /**

compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts

@@ -25,7 +25,6 @@ import {
   makeBlockId,
   makeInstructionId,
   makePropertyLiteral,
-  makeType,
   markInstructionIds,
   promoteTemporary,
   reversePostorderBlocks,
@@ -253,7 +252,6 @@ function emitSelectorFn(env: Environment, keys: Array<string>): Instruction {
     env,
     params: [obj],
     returnTypeAnnotation: null,
-    returnType: makeType(),
     returns: createTemporaryPlace(env, GeneratedSource),
     context: [],
     effects: null,

compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts

@@ -21,7 +21,6 @@ import {
   makeBlockId,
   makeIdentifierName,
   makeInstructionId,
-  makeType,
   ObjectProperty,
   Place,
   promoteTemporary,
@@ -368,7 +367,6 @@ function emitOutlinedFn(
     env,
     params: [propsObj],
     returnTypeAnnotation: null,
-    returnType: makeType(),
     returns: createTemporaryPlace(env, GeneratedSource),
     context: [],
     effects: null,

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -349,11 +349,9 @@ function codegenReactiveFunction(
   fn: ReactiveFunction,
 ): Result<CodegenFunction, CompilerError> {
   for (const param of fn.params) {
-    if (param.kind === 'Identifier') {
-      cx.temp.set(param.identifier.declarationId, null);
-    } else {
-      cx.temp.set(param.place.identifier.declarationId, null);
-    }
+    const place = param.kind === 'Identifier' ? param : param.place;
+    cx.temp.set(place.identifier.declarationId, null);
+    cx.declare(place.identifier);
   }
 
   const params = fn.params.map(param => convertParameter(param));
@@ -1183,7 +1181,7 @@ function codegenTerminal(
               ? codegenPlaceToExpression(cx, case_.test)
               : null;
           const block = codegenBlock(cx, case_.block!);
-          return t.switchCase(test, [block]);
+          return t.switchCase(test, block.body.length === 0 ? [] : [block]);
         }),
       );
     }

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/ExtractScopeDeclarationsFromDestructuring.ts

@@ -79,6 +79,10 @@ export function extractScopeDeclarationsFromDestructuring(
   fn: ReactiveFunction,
 ): void {
   const state = new State(fn.env);
+  for (const param of fn.params) {
+    const place = param.kind === 'Identifier' ? param : param.place;
+    state.declared.add(place.identifier.declarationId);
+  }
   visitReactiveFunction(fn, new Visitor(), state);
 }
 

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

@@ -90,7 +90,8 @@ function apply(func: HIRFunction, unifier: Unifier): void {
       }
     }
   }
-  func.returnType = unifier.get(func.returnType);
+  const returns = func.returns.identifier;
+  returns.type = unifier.get(returns.type);
 }
 
 type TypeEquation = {
@@ -143,12 +144,12 @@ function* generate(
     }
   }
   if (returnTypes.length > 1) {
-    yield equation(func.returnType, {
+    yield equation(func.returns.identifier.type, {
       kind: 'Phi',
       operands: returnTypes,
     });
   } else if (returnTypes.length === 1) {
-    yield equation(func.returnType, returnTypes[0]!);
+    yield equation(func.returns.identifier.type, returnTypes[0]!);
   }
 }
 
@@ -407,7 +408,7 @@ function* generateInstructionTypes(
       yield equation(left, {
         kind: 'Function',
         shapeId: BuiltInFunctionId,
-        return: value.loweredFunc.func.returnType,
+        return: value.loweredFunc.func.returns.identifier.type,
         isConstructor: false,
       });
       break;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/block-scoping-switch-variable-scoping.expect.md

@@ -50,8 +50,7 @@ function Component(props) {
       console.log(handlers.value);
       break bb0;
     }
-    default: {
-    }
+    default:
   }
 
   t0 = handlers;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/dominator.expect.md

@@ -67,8 +67,7 @@ function Component(props) {
     case "b": {
       break bb1;
     }
-    case "c": {
-    }
+    case "c":
     default: {
       x = 6;
     }

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

@@ -24,8 +24,6 @@ export const FIXTURE_ENTRYPOINT = {
   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.invalid-access-ref-during-render.expect.md

@@ -19,6 +19,8 @@ function Component(props) {
   3 |   const ref = useRef(null);
 > 4 |   const value = ref.current;
     |                 ^^^^^^^^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (4:4)
+
+InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef). Cannot access ref value `value` (5:5)
   5 |   return value;
   6 | }
   7 |

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-aliased-ref-in-callback-invoked-during-render-.expect.md

@@ -19,12 +19,17 @@ function Component(props) {
 ## Error
 
 ```
-   7 |     return <Foo item={item} current={current} />;
-   8 |   };
->  9 |   return <Items>{props.items.map(item => renderItem(item))}</Items>;
-     |                                  ^^^^^^^^^^^^^^^^^^^^^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (9:9)
-  10 | }
-  11 |
+  4 |   const renderItem = item => {
+  5 |     const aliasedRef = ref;
+> 6 |     const current = aliasedRef.current;
+    |                     ^^^^^^^^^^^^^^^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (6:6)
+
+InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef). Cannot access ref value `current` (7:7)
+
+InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (7:7)
+  7 |     return <Foo item={item} current={current} />;
+  8 |   };
+  9 |   return <Items>{props.items.map(item => renderItem(item))}</Items>;
 ```