[compiler] ref guards apply up to fallthrough of the test (#34024)

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.

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/34024).
* #34027
* #34026
* #34025
* __->__ #34024

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