[compiler] disallow ref access in state initializer, reducer/initializer

Per title, disallow ref access in `useState()` initializer function, `useReducer()` reducer, and `useReducer()` init function.

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

@@ -80,18 +80,8 @@ type RefAccessRefType =
 
 type RefFnType = {readRefEffect: boolean; returnType: RefAccessType};
 
-class Env {
+class Env extends Map<IdentifierId, RefAccessType> {
   #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;
@@ -101,14 +91,8 @@ class Env {
     return this.#changed;
   }
 
-  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);
+  override set(key: IdentifierId, value: RefAccessType): this {
+    const cur = this.get(key);
     const widenedValue = joinRefAccessTypes(value, cur ?? {kind: 'None'});
     if (
       !(cur == null && widenedValue.kind === 'None') &&
@@ -116,8 +100,7 @@ class Env {
     ) {
       this.#changed = true;
     }
-    this.#data.set(operandId, widenedValue);
-    return this;
+    return super.set(key, widenedValue);
   }
 }
 
@@ -125,48 +108,9 @@ 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'};
@@ -580,25 +524,11 @@ function validateNoRefAccessInRenderImpl(
             } else {
               validateNoRefUpdate(errors, env, instr.value.object, instr.loc);
             }
-            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);
+            for (const operand of eachInstructionValueOperand(instr.value)) {
+              if (operand === instr.value.object) {
+                continue;
               }
+              validateNoRefValueAccess(errors, env, operand);
             }
             break;
           }
@@ -800,7 +730,11 @@ function validateNoRefUpdate(
   loc: SourceLocation,
 ): void {
   const type = destructure(env.get(operand.identifier.id));
-  if (type?.kind === 'Ref' || type?.kind === 'RefValue') {
+  if (
+    type?.kind === 'Ref' ||
+    type?.kind === 'RefValue' ||
+    (type?.kind === 'Structure' && type.fn?.readRefEffect)
+  ) {
     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

@@ -1,52 +0,0 @@
-
-## 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

@@ -1,14 +0,0 @@
-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

@@ -1,37 +0,0 @@
-
-## 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

@@ -1,9 +0,0 @@
-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,13 +41,14 @@ 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:12:28
-  10 |   const x = {a, val: val.ref.current};
+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
   11 |
-> 12 |   return <VideoList videos={x} />;
-     |                             ^ Cannot access ref value during render
+  12 |   return <VideoList videos={x} />;
   13 | }
-  14 |
 ```