[compiler] ref guards apply up to fallthrough of the test

Fixes #30782

When developers do an `if (ref.current == null)` guard for lazy ref initialization, the "safe" blocks should extend up to the if's fallthrough. Previously we only allowed writing to the ref in the if consequent, but this meant that you couldn't use a ternary, logical, etc in the if body.

Author: josephsavona

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

@@ -27,6 +27,7 @@ import {
   eachTerminalOperand,
 } from '../HIR/visitors';
 import {Err, Ok, Result} from '../Utils/Result';
+import {retainWhere} from '../Utils/utils';
 
 /**
  * Validates that a function does not access a ref value during render. This includes a partial check
@@ -279,9 +280,10 @@ function validateNoRefAccessInRenderImpl(
   for (let i = 0; (i == 0 || env.hasChanged()) && i < 10; i++) {
     env.resetChanged();
     returnValues = [];
-    const safeBlocks = new Map<BlockId, RefId>();
+    const safeBlocks: Array<{block: BlockId; ref: RefId}> = [];
     const errors = new CompilerError();
     for (const [, block] of fn.body.blocks) {
+      retainWhere(safeBlocks, entry => entry.block !== block.id);
       for (const phi of block.phis) {
         env.set(
           phi.place.identifier.id,
@@ -503,15 +505,17 @@ function validateNoRefAccessInRenderImpl(
           case 'PropertyStore':
           case 'ComputedDelete':
           case 'ComputedStore': {
-            const safe = safeBlocks.get(block.id);
             const target = env.get(instr.value.object.identifier.id);
+            let safe: (typeof safeBlocks)['0'] | null | undefined = null;
             if (
               instr.value.kind === 'PropertyStore' &&
-              safe != null &&
-              target?.kind === 'Ref' &&
-              target.refId === safe
+              target != null &&
+              target.kind === 'Ref'
             ) {
-              safeBlocks.delete(block.id);
+              safe = safeBlocks.find(entry => entry.ref === target.refId);
+            }
+            if (safe != null) {
+              retainWhere(safeBlocks, entry => entry !== safe);
             } else {
               validateNoRefUpdate(errors, env, instr.value.object, instr.loc);
             }
@@ -599,8 +603,11 @@ function validateNoRefAccessInRenderImpl(
 
       if (block.terminal.kind === 'if') {
         const test = env.get(block.terminal.test.identifier.id);
-        if (test?.kind === 'Guard') {
-          safeBlocks.set(block.terminal.consequent, test.refId);
+        if (
+          test?.kind === 'Guard' &&
+          safeBlocks.find(entry => entry.ref === test.refId) == null
+        ) {
+          safeBlocks.push({block: block.terminal.fallthrough, ref: test.refId});
         }
       }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-lazy-initialization-with-logical.expect.md

@@ -0,0 +1,68 @@
+
+## Input
+
+```javascript
+// @validateRefAccessDuringRender
+
+import {useRef} from 'react';
+
+function Component(props) {
+  const ref = useRef(null);
+  if (ref.current == null) {
+    // the logical means the ref write is in a different block
+    // from the if consequent. this tests that the "safe" blocks
+    // extend up to the if's fallthrough
+    ref.current = props.unknownKey ?? props.value;
+  }
+  return <Child ref={ref} />;
+}
+
+function Child({ref}) {
+  'use no memo';
+  return ref.current;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{value: 42}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validateRefAccessDuringRender
+
+import { useRef } from "react";
+
+function Component(props) {
+  const $ = _c(1);
+  const ref = useRef(null);
+  if (ref.current == null) {
+    ref.current = props.unknownKey ?? props.value;
+  }
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <Child ref={ref} />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+function Child({ ref }) {
+  "use no memo";
+  return ref.current;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) 42

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-ref-lazy-initialization-with-logical.js

@@ -0,0 +1,24 @@
+// @validateRefAccessDuringRender
+
+import {useRef} from 'react';
+
+function Component(props) {
+  const ref = useRef(null);
+  if (ref.current == null) {
+    // the logical means the ref write is in a different block
+    // from the if consequent. this tests that the "safe" blocks
+    // extend up to the if's fallthrough
+    ref.current = props.unknownKey ?? props.value;
+  }
+  return <Child ref={ref} />;
+}
+
+function Child({ref}) {
+  'use no memo';
+  return ref.current;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{value: 42}],
+};