[compiler] Refactor event handler validation to use type inference

Author: kolvian

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

@@ -670,6 +670,15 @@ export const EnvironmentConfigSchema = z.object({
    * from refs need to be stored in state during mount.
    */
   enableAllowSetStateFromRefsInEffects: z.boolean().default(true),
+
+  /**
+   * Enables inference of event handler types for JSX props on built-in DOM elements.
+   * When enabled, functions passed to event handler props (props starting with "on")
+   * on primitive JSX tags are inferred to have the BuiltinEventHandlerId type, which
+   * allows ref access within those functions since DOM event handlers are guaranteed
+   * by React to only execute in response to events, not during render.
+   */
+  enableInferEventHandlers: z.boolean().default(false),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;

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

@@ -29,7 +29,7 @@ import {
   BuiltInUseTransitionId,
   BuiltInWeakMapId,
   BuiltInWeakSetId,
-  BuiltinEffectEventId,
+  BuiltInEffectEventId,
   ReanimatedSharedValueId,
   ShapeRegistry,
   addFunction,
@@ -863,7 +863,7 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
         returnType: {
           kind: 'Function',
           return: {kind: 'Poly'},
-          shapeId: BuiltinEffectEventId,
+          shapeId: BuiltInEffectEventId,
           isConstructor: false,
         },
         calleeEffect: Effect.Read,

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

@@ -403,8 +403,9 @@ export const BuiltInStartTransitionId = 'BuiltInStartTransition';
 export const BuiltInFireId = 'BuiltInFire';
 export const BuiltInFireFunctionId = 'BuiltInFireFunction';
 export const BuiltInUseEffectEventId = 'BuiltInUseEffectEvent';
-export const BuiltinEffectEventId = 'BuiltInEffectEventFunction';
+export const BuiltInEffectEventId = 'BuiltInEffectEventFunction';
 export const BuiltInAutodepsId = 'BuiltInAutoDepsId';
+export const BuiltInEventHandlerId = 'BuiltInEventHandlerId';
 
 // See getReanimatedModuleType() in Globals.ts — this is part of supporting Reanimated's ref-like types
 export const ReanimatedSharedValueId = 'ReanimatedSharedValueId';
@@ -1243,7 +1244,20 @@ addFunction(
     calleeEffect: Effect.ConditionallyMutate,
     returnValueKind: ValueKind.Mutable,
   },
-  BuiltinEffectEventId,
+  BuiltInEffectEventId,
+);
+
+addFunction(
+  BUILTIN_SHAPES,
+  [],
+  {
+    positionalParams: [],
+    restParam: Effect.ConditionallyMutate,
+    returnType: {kind: 'Poly'},
+    calleeEffect: Effect.ConditionallyMutate,
+    returnValueKind: ValueKind.Mutable,
+  },
+  BuiltInEventHandlerId,
 );
 
 /**

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

@@ -25,6 +25,7 @@ import {
 } from '../HIR/HIR';
 import {
   BuiltInArrayId,
+  BuiltInEventHandlerId,
   BuiltInFunctionId,
   BuiltInJsxId,
   BuiltInMixedReadonlyId,
@@ -471,6 +472,32 @@ function* generateInstructionTypes(
           }
         }
       }
+      if (env.config.enableInferEventHandlers) {
+        if (value.kind === 'JsxExpression' && value.tag.kind === 'BuiltinTag') {
+          /*
+           * Infer event handler types for built-in DOM elements.
+           * Props starting with "on" (e.g., onClick, onSubmit) on primitive tags
+           * are inferred as event handlers. This allows functions with ref access
+           * to be passed to these props, since DOM event handlers are guaranteed
+           * by React to only execute in response to events, never during render.
+           */
+          for (const prop of value.props) {
+            if (
+              prop.kind === 'JsxAttribute' &&
+              prop.name.startsWith('on') &&
+              prop.name.length > 2 &&
+              prop.name[2] === prop.name[2].toUpperCase()
+            ) {
+              yield equation(prop.place.identifier.type, {
+                kind: 'Function',
+                shapeId: BuiltInEventHandlerId,
+                return: makeType(),
+                isConstructor: false,
+              });
+            }
+          }
+        }
+      }
       yield equation(left, {kind: 'Object', shapeId: BuiltInJsxId});
       break;
     }

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

@@ -14,12 +14,14 @@ import {
   BlockId,
   HIRFunction,
   IdentifierId,
+  Identifier,
   Place,
   SourceLocation,
   getHookKindForType,
   isRefValueType,
   isUseRefType,
 } from '../HIR';
+import {BuiltInEventHandlerId} from '../HIR/ObjectShape';
 import {
   eachInstructionOperand,
   eachInstructionValueOperand,
@@ -183,6 +185,11 @@ function refTypeOfType(place: Place): RefAccessType {
   }
 }
 
+function isEventHandlerType(identifier: Identifier): boolean {
+  const type = identifier.type;
+  return type.kind === 'Function' && type.shapeId === BuiltInEventHandlerId;
+}
+
 function tyEqual(a: RefAccessType, b: RefAccessType): boolean {
   if (a.kind !== b.kind) {
     return false;
@@ -354,7 +361,6 @@ function validateNoRefAccessInRenderImpl(
   }
 
   const interpolatedAsJsx = new Set<IdentifierId>();
-  const usedAsEventHandlerProp = new Set<IdentifierId>();
   for (const block of fn.body.blocks.values()) {
     for (const instr of block.instructions) {
       const {value} = instr;
@@ -364,27 +370,6 @@ function validateNoRefAccessInRenderImpl(
             interpolatedAsJsx.add(child.identifier.id);
           }
         }
-        if (value.kind === 'JsxExpression') {
-          /*
-           * Track identifiers used as event handler props on built-in DOM elements.
-           * We only allow this optimization for native DOM elements because their
-           * event handlers are guaranteed by React to only execute in response to
-           * events, never during render. Custom components could technically call
-           * their onFoo props during render, which would violate ref access rules.
-           */
-          if (value.tag.kind === 'BuiltinTag') {
-            for (const prop of value.props) {
-              if (
-                prop.kind === 'JsxAttribute' &&
-                prop.name.startsWith('on') &&
-                prop.name.length > 2 &&
-                prop.name[2] === prop.name[2].toUpperCase()
-              ) {
-                usedAsEventHandlerProp.add(prop.place.identifier.id);
-              }
-            }
-          }
-        }
       }
     }
   }
@@ -541,8 +526,8 @@ function validateNoRefAccessInRenderImpl(
              */
             if (!didError) {
               const isRefLValue = isUseRefType(instr.lvalue.identifier);
-              const isUsedAsEventHandler = usedAsEventHandlerProp.has(
-                instr.lvalue.identifier.id,
+              const isEventHandlerLValue = isEventHandlerType(
+                instr.lvalue.identifier,
               );
               for (const operand of eachInstructionValueOperand(instr.value)) {
                 /**
@@ -551,29 +536,16 @@ function validateNoRefAccessInRenderImpl(
                  */
                 if (
                   isRefLValue ||
+                  isEventHandlerLValue ||
                   (hookKind != null &&
                     hookKind !== 'useState' &&
                     hookKind !== 'useReducer')
                 ) {
                   /**
-                   * Special cases:
-                   *
-                   * 1. the lvalue is a ref
-                   * In general passing a ref to a function may access that ref
-                   * value during render, so we disallow it.
-                   *
-                   * The main exception is the "mergeRefs" pattern, ie a function
-                   * that accepts multiple refs as arguments (or an array of refs)
-                   * and returns a new, aggregated ref. If the lvalue is a ref,
-                   * we assume that the user is doing this pattern and allow passing
-                   * refs.
-                   *
-                   * Eg `const mergedRef = mergeRefs(ref1, ref2)`
-                   *
-                   * 2. calling hooks
-                   *
-                   * Hooks are independently checked to ensure they don't access refs
-                   * during render.
+                   * Allow passing refs or ref-accessing functions when:
+                   * 1. lvalue is a ref (mergeRefs pattern: `mergeRefs(ref1, ref2)`)
+                   * 2. lvalue is an event handler (DOM events execute outside render)
+                   * 3. calling hooks (independently validated for ref safety)
                    */
                   validateNoDirectRefValueAccess(errors, operand, env);
                 } else if (interpolatedAsJsx.has(instr.lvalue.identifier.id)) {
@@ -585,25 +557,6 @@ function validateNoRefAccessInRenderImpl(
                    * render function which attempts to obey the rules.
                    */
                   validateNoRefValueAccess(errors, env, operand);
-                } else if (isUsedAsEventHandler) {
-                  /**
-                   * Special case: the lvalue is used as an event handler prop
-                   * on a built-in DOM element
-                   *
-                   * For example `<form onSubmit={handleSubmit(onSubmit)}>`. Here
-                   * handleSubmit is wrapping onSubmit to create an event handler.
-                   * Functions passed to event handler wrappers can safely access
-                   * refs because built-in DOM event handlers are guaranteed by React
-                   * to only execute in response to actual events, never during render.
-                   *
-                   * We only allow this for built-in DOM elements (not custom components)
-                   * because custom components could technically call their onFoo props
-                   * during render, which would violate ref access rules.
-                   *
-                   * We allow passing functions with ref access to these wrappers,
-                   * but still validate that direct ref values aren't passed.
-                   */
-                  validateNoDirectRefValueAccess(errors, operand, env);
                 } else {
                   validateNoRefPassedToFunction(
                     errors,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-async-event-handler-wrapper.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @enableInferEventHandlers
 // @validateRefAccessDuringRender
 import {useRef} from 'react';
 
@@ -56,7 +57,8 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @validateRefAccessDuringRender
+import { c as _c } from "react/compiler-runtime"; // @enableInferEventHandlers
+// @validateRefAccessDuringRender
 import { useRef } from "react";
 
 // Simulates react-hook-form's handleSubmit

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-async-event-handler-wrapper.tsx

@@ -1,3 +1,4 @@
+// @enableInferEventHandlers
 // @validateRefAccessDuringRender
 import {useRef} from 'react';
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-event-handler-wrapper.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @enableInferEventHandlers
 // @validateRefAccessDuringRender
 import {useRef} from 'react';
 
@@ -44,7 +45,8 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { c as _c } from "react/compiler-runtime"; // @validateRefAccessDuringRender
+import { c as _c } from "react/compiler-runtime"; // @enableInferEventHandlers
+// @validateRefAccessDuringRender
 import { useRef } from "react";
 
 // Simulates react-hook-form's handleSubmit or similar event handler wrappers

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-access-in-event-handler-wrapper.tsx

@@ -1,3 +1,4 @@
+// @enableInferEventHandlers
 // @validateRefAccessDuringRender
 import {useRef} from 'react';
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.ref-value-in-custom-component-event-handler-wrapper.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+// @enableInferEventHandlers
 // @validateRefAccessDuringRender
 import {useRef} from 'react';
 
@@ -56,14 +57,14 @@ Error: Cannot access refs during render
 
 React refs are values that are not needed for rendering. Refs should only be accessed outside of render, such as in event handlers or effects. Accessing a ref value (the `current` property) during render can cause your component not to update as expected (https://react.dev/reference/react/useRef).
 
-error.ref-value-in-custom-component-event-handler-wrapper.ts:31:41
-  29 |     <>
-  30 |       <input ref={ref} />
-> 31 |       <CustomForm onSubmit={handleSubmit(onSubmit)}>
+error.ref-value-in-custom-component-event-handler-wrapper.ts:32:41
+  30 |     <>
+  31 |       <input ref={ref} />
+> 32 |       <CustomForm onSubmit={handleSubmit(onSubmit)}>
      |                                          ^^^^^^^^ Passing a ref to a function may read its value during render
-  32 |         <button type="submit">Submit</button>
-  33 |       </CustomForm>
-  34 |     </>
+  33 |         <button type="submit">Submit</button>
+  34 |       </CustomForm>
+  35 |     </>
 ```