[compiler] Allow passing refs to render helpers

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

@@ -451,6 +451,18 @@ function* generateInstructionTypes(
 
     case 'JsxExpression':
     case 'JsxFragment': {
+      if (env.config.enableTreatRefLikeIdentifiersAsRefs) {
+        if (value.kind === 'JsxExpression') {
+          for (const prop of value.props) {
+            if (prop.kind === 'JsxAttribute' && prop.name === 'ref') {
+              yield equation(prop.place.identifier.type, {
+                kind: 'Object',
+                shapeId: BuiltInUseRefId,
+              });
+            }
+          }
+        }
+      }
       yield equation(left, {kind: 'Object', shapeId: BuiltInJsxId});
       break;
     }

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 = [];
@@ -407,13 +421,48 @@ function validateNoRefAccessInRenderImpl(
                 );
               }
             }
+            /*
+             * If we already reported an error on this instruction, don't report
+             * duplicate errors
+             */
             if (!didError) {
-              /*
-               * If we already reported an error on this instruction, don't report
-               * duplicate errors
-               */
+              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 {
                   validateNoRefPassedToFunction(

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-merge-refs-pattern.expect.md

@@ -0,0 +1,44 @@
+
+## Input
+
+```javascript
+// @enableTreatRefLikeIdentifiersAsRefs @validateRefAccessDuringRender
+
+import {useRef} from 'react';
+
+function Component() {
+  const ref = useRef(null);
+  const ref2 = useRef(null);
+  const mergedRef = mergeRefs([ref], ref2);
+
+  return <Stringify ref={mergedRef} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableTreatRefLikeIdentifiersAsRefs @validateRefAccessDuringRender
+
+import { useRef } from "react";
+
+function Component() {
+  const $ = _c(1);
+  const ref = useRef(null);
+  const ref2 = useRef(null);
+  const mergedRef = mergeRefs([ref], ref2);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = <Stringify ref={mergedRef} />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+```
+
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-merge-refs-pattern.js

@@ -0,0 +1,11 @@
+// @enableTreatRefLikeIdentifiersAsRefs @validateRefAccessDuringRender
+
+import {useRef} from 'react';
+
+function Component() {
+  const ref = useRef(null);
+  const ref2 = useRef(null);
+  const mergedRef = mergeRefs([ref], ref2);
+
+  return <Stringify ref={mergedRef} />;
+}

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