[compiler] Allow passing refs to render helpers (#34006)

We infer render helpers as functions whose result is immediately
interpolated into jsx. This is a very conservative approximation, to
help with common cases like `<Foo>{props.renderItem(ref)}</Foo>`. The
idea is similar to hooks that it's ultimately on the developer to catch
ref-in-render validations (and the runtime detects them too), so we can
be a bit more relaxed since there are valid reasons to use this pattern.

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

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

@@ -262,6 +262,20 @@ function validateNoRefAccessInRenderImpl(
     env.set(place.identifier.id, type);
   }
 
+  const interpolatedAsJsx = new Set<IdentifierId>();
+  for (const block of fn.body.blocks.values()) {
+    for (const instr of block.instructions) {
+      const {value} = instr;
+      if (value.kind === 'JsxExpression' || value.kind === 'JsxFragment') {
+        if (value.children != null) {
+          for (const child of value.children) {
+            interpolatedAsJsx.add(child.identifier.id);
+          }
+        }
+      }
+    }
+  }
+
   for (let i = 0; (i == 0 || env.hasChanged()) && i < 10; i++) {
     env.resetChanged();
     returnValues = [];
@@ -414,7 +428,41 @@ function validateNoRefAccessInRenderImpl(
             if (!didError) {
               const isRefLValue = isUseRefType(instr.lvalue.identifier);
               for (const operand of eachInstructionValueOperand(instr.value)) {
-                if (hookKind != null) {
+                /**
+                 * By default we check that function call operands are not refs,
+                 * ref values, or functions that can access refs.
+                 */
+                if (
+                  isRefLValue ||
+                  interpolatedAsJsx.has(instr.lvalue.identifier.id) ||
+                  hookKind != null
+                ) {
+                  /**
+                   * 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) the lvalue is passed as a jsx child
+                   *
+                   * For example `<Foo>{renderHelper(ref)}</Foo>`. Here we have more
+                   * context and infer that the ref is being passed to a component-like
+                   * render function which attempts to obey the rules.
+                   *
+                   * 3) hooks
+                   *
+                   * Hooks are independently checked to ensure they don't access refs
+                   * during render.
+                   */
                   validateNoDirectRefValueAccess(errors, operand, env);
                 } else if (!isRefLValue) {
                   /**

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-passing-ref-to-render-helper-props-object.expect.md

@@ -0,0 +1,45 @@
+
+## Input
+
+```javascript
+// @enableTreatRefLikeIdentifiersAsRefs @validateRefAccessDuringRender
+
+import {useRef} from 'react';
+
+function Component(props) {
+  const ref = useRef(null);
+
+  return <Foo>{props.render({ref})}</Foo>;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableTreatRefLikeIdentifiersAsRefs @validateRefAccessDuringRender
+
+import { useRef } from "react";
+
+function Component(props) {
+  const $ = _c(3);
+  const ref = useRef(null);
+
+  const T0 = Foo;
+  const t0 = props.render({ ref });
+  let t1;
+  if ($[0] !== T0 || $[1] !== t0) {
+    t1 = <T0>{t0}</T0>;
+    $[0] = T0;
+    $[1] = t0;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-passing-ref-to-render-helper-props-object.js

@@ -0,0 +1,9 @@
+// @enableTreatRefLikeIdentifiersAsRefs @validateRefAccessDuringRender
+
+import {useRef} from 'react';
+
+function Component(props) {
+  const ref = useRef(null);
+
+  return <Foo>{props.render({ref})}</Foo>;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-passing-ref-to-render-helper.expect.md

@@ -0,0 +1,49 @@
+
+## Input
+
+```javascript
+// @enableTreatRefLikeIdentifiersAsRefs @validateRefAccessDuringRender
+
+import {useRef} from 'react';
+
+function Component(props) {
+  const ref = useRef(null);
+
+  return <Foo>{props.render(ref)}</Foo>;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableTreatRefLikeIdentifiersAsRefs @validateRefAccessDuringRender
+
+import { useRef } from "react";
+
+function Component(props) {
+  const $ = _c(4);
+  const ref = useRef(null);
+  let t0;
+  if ($[0] !== props.render) {
+    t0 = props.render(ref);
+    $[0] = props.render;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  let t1;
+  if ($[2] !== t0) {
+    t1 = <Foo>{t0}</Foo>;
+    $[2] = t0;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-passing-ref-to-render-helper.js

@@ -0,0 +1,9 @@
+// @enableTreatRefLikeIdentifiersAsRefs @validateRefAccessDuringRender
+
+import {useRef} from 'react';
+
+function Component(props) {
+  const ref = useRef(null);
+
+  return <Foo>{props.render(ref)}</Foo>;
+}