[compiler] Allow assigning ref-accessing functions to objects if not mutated (facebook#34026)

Allows assigning a ref-accessing function to an object so long as that
object is not subsequently transitively mutated. We should likely
rewrite the ref validation to use the new mutation/aliasing effects,
which would provide a more consistent behavior across instruction types
and require fewer special cases like this.

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

Author: josephsavona

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

@@ -80,8 +80,18 @@ type RefAccessRefType =
 
 type RefFnType = {readRefEffect: boolean; returnType: RefAccessType};
 
-class Env extends Map<IdentifierId, RefAccessType> {
+class Env {
   #changed = false;
+  #data: Map<IdentifierId, RefAccessType> = new Map();
+  #temporaries: Map<IdentifierId, Place> = new Map();
+
+  lookup(place: Place): Place {
+    return this.#temporaries.get(place.identifier.id) ?? place;
+  }
+
+  define(place: Place, value: Place): void {
+    this.#temporaries.set(place.identifier.id, value);
+  }
 
   resetChanged(): void {
     this.#changed = false;
@@ -91,26 +101,72 @@ class Env extends Map<IdentifierId, RefAccessType> {
     return this.#changed;
   }
 
-  override set(key: IdentifierId, value: RefAccessType): this {
-    const cur = this.get(key);
+  get(key: IdentifierId): RefAccessType | undefined {
+    const operandId = this.#temporaries.get(key)?.identifier.id ?? key;
+    return this.#data.get(operandId);
+  }
+
+  set(key: IdentifierId, value: RefAccessType): this {
+    const operandId = this.#temporaries.get(key)?.identifier.id ?? key;
+    const cur = this.#data.get(operandId);
     const widenedValue = joinRefAccessTypes(value, cur ?? {kind: 'None'});
     if (
       !(cur == null && widenedValue.kind === 'None') &&
       (cur == null || !tyEqual(cur, widenedValue))
     ) {
       this.#changed = true;
     }
-    return super.set(key, widenedValue);
+    this.#data.set(operandId, widenedValue);
+    return this;
   }
 }
 
 export function validateNoRefAccessInRender(
   fn: HIRFunction,
 ): Result<void, CompilerError> {
   const env = new Env();
+  collectTemporariesSidemap(fn, env);
   return validateNoRefAccessInRenderImpl(fn, env).map(_ => undefined);
 }
 
+function collectTemporariesSidemap(fn: HIRFunction, env: Env): void {
+  for (const block of fn.body.blocks.values()) {
+    for (const instr of block.instructions) {
+      const {lvalue, value} = instr;
+      switch (value.kind) {
+        case 'LoadLocal': {
+          const temp = env.lookup(value.place);
+          if (temp != null) {
+            env.define(lvalue, temp);
+          }
+          break;
+        }
+        case 'StoreLocal': {
+          const temp = env.lookup(value.value);
+          if (temp != null) {
+            env.define(lvalue, temp);
+            env.define(value.lvalue.place, temp);
+          }
+          break;
+        }
+        case 'PropertyLoad': {
+          if (
+            isUseRefType(value.object.identifier) &&
+            value.property === 'current'
+          ) {
+            continue;
+          }
+          const temp = env.lookup(value.object);
+          if (temp != null) {
+            env.define(lvalue, temp);
+          }
+          break;
+        }
+      }
+    }
+  }
+}
+
 function refTypeOfType(place: Place): RefAccessType {
   if (isRefValueType(place.identifier)) {
     return {kind: 'RefValue'};
@@ -524,11 +580,25 @@ function validateNoRefAccessInRenderImpl(
             } else {
               validateNoRefUpdate(errors, env, instr.value.object, instr.loc);
             }
-            for (const operand of eachInstructionValueOperand(instr.value)) {
-              if (operand === instr.value.object) {
-                continue;
+            if (
+              instr.value.kind === 'ComputedDelete' ||
+              instr.value.kind === 'ComputedStore'
+            ) {
+              validateNoRefValueAccess(errors, env, instr.value.property);
+            }
+            if (
+              instr.value.kind === 'ComputedStore' ||
+              instr.value.kind === 'PropertyStore'
+            ) {
+              validateNoDirectRefValueAccess(errors, instr.value.value, env);
+              const type = env.get(instr.value.value.identifier.id);
+              if (type != null && type.kind === 'Structure') {
+                let objectType: RefAccessType = type;
+                if (target != null) {
+                  objectType = joinRefAccessTypes(objectType, target);
+                }
+                env.set(instr.value.object.identifier.id, objectType);
               }
-              validateNoRefValueAccess(errors, env, operand);
             }
             break;
           }
@@ -730,11 +800,7 @@ function validateNoRefUpdate(
   loc: SourceLocation,
 ): void {
   const type = destructure(env.get(operand.identifier.id));
-  if (
-    type?.kind === 'Ref' ||
-    type?.kind === 'RefValue' ||
-    (type?.kind === 'Structure' && type.fn?.readRefEffect)
-  ) {
+  if (type?.kind === 'Ref' || type?.kind === 'RefValue') {
     errors.pushDiagnostic(
       CompilerDiagnostic.create({
         severity: ErrorSeverity.InvalidReact,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-assigning-ref-accessing-function-to-object-property-if-not-mutated.expect.md

@@ -0,0 +1,52 @@
+
+## Input
+
+```javascript
+import {useRef} from 'react';
+import {Stringify} from 'shared-runtime';
+
+function Component(props) {
+  const ref = useRef(props.value);
+  const object = {};
+  object.foo = () => ref.current;
+  return <Stringify object={object} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{value: 42}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useRef } from "react";
+import { Stringify } from "shared-runtime";
+
+function Component(props) {
+  const $ = _c(1);
+  const ref = useRef(props.value);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const object = {};
+    object.foo = () => ref.current;
+    t0 = <Stringify object={object} shouldInvokeFns={true} />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"object":{"foo":{"kind":"Function","result":42}},"shouldInvokeFns":true}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-assigning-ref-accessing-function-to-object-property-if-not-mutated.js

@@ -0,0 +1,14 @@
+import {useRef} from 'react';
+import {Stringify} from 'shared-runtime';
+
+function Component(props) {
+  const ref = useRef(props.value);
+  const object = {};
+  object.foo = () => ref.current;
+  return <Stringify object={object} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{value: 42}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-access-ref-in-render-mutate-object-with-ref-function.expect.md

@@ -0,0 +1,37 @@
+
+## Input
+
+```javascript
+import {useRef} from 'react';
+
+function Component() {
+  const ref = useRef(null);
+  const object = {};
+  object.foo = () => ref.current;
+  const refValue = object.foo();
+  return <div>{refValue}</div>;
+}
+
+```
+
+
+## Error
+
+```
+Found 1 error:
+
+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.invalid-access-ref-in-render-mutate-object-with-ref-function.ts:7:19
+   5 |   const object = {};
+   6 |   object.foo = () => ref.current;
+>  7 |   const refValue = object.foo();
+     |                    ^^^^^^^^^^ This function accesses a ref value
+   8 |   return <div>{refValue}</div>;
+   9 | }
+  10 |
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-access-ref-in-render-mutate-object-with-ref-function.js

@@ -0,0 +1,9 @@
+import {useRef} from 'react';
+
+function Component() {
+  const ref = useRef(null);
+  const object = {};
+  object.foo = () => ref.current;
+  const refValue = object.foo();
+  return <div>{refValue}</div>;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-use-ref-added-to-dep-without-type-info.expect.md

@@ -41,14 +41,13 @@ 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.invalid-use-ref-added-to-dep-without-type-info.ts:10:21
-   8 |   // however, this is an instance of accessing a ref during render and is disallowed
-   9 |   // under React's rules, so we reject this input
-> 10 |   const x = {a, val: val.ref.current};
-     |                      ^^^^^^^^^^^^^^^ Cannot access ref value during render
+error.invalid-use-ref-added-to-dep-without-type-info.ts:12:28
+  10 |   const x = {a, val: val.ref.current};
   11 |
-  12 |   return <VideoList videos={x} />;
+> 12 |   return <VideoList videos={x} />;
+     |                             ^ Cannot access ref value during render
   13 | }
+  14 |
 ```