[compiler] Exclude refs and ref values from having mutable ranges

Summary: Refs, as stable values that the rules of react around mutability do not apply to, currently are treated as having mutable ranges, and through aliasing, this can extend the mutable range for other values and disrupt good memoization for those values. This PR excludes refs and their .current values from having mutable ranges. Note that this is unsafe if ref access is allowed in render: if a mutable value is assigned to ref.current and then ref.current is mutated later, we won't realize that the original mutable value's range extends. ghstack-source-id: e8f36ac25e2c9aadb0bf13bd8142e4593ee9f984 Pull Request resolved: #30713
facebook · Aug 16, 2024 · 5030e08 · 5030e08
1 parent 85fb95c
commit 5030e08
Show file tree

Hide file tree

Showing 17 changed files with 159 additions and 130 deletions.
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
@@ -1591,6 +1591,10 @@ export function isUseStateType(id: Identifier): boolean {
   return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInUseState';
 }
 
+export function isRefOrRefValue(id: Identifier): boolean {
+  return isUseRefType(id) || isRefValueType(id);
+}
+
 export function isSetStateType(id: Identifier): boolean {
   return id.type.kind === 'Function' && id.type.shapeId === 'BuiltInSetState';
 }

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/AnalyseFunctions.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/AnalyseFunctions.ts
@@ -14,8 +14,7 @@ import {
   LoweredFunction,
   Place,
   ReactiveScopeDependency,
-  isRefValueType,
-  isUseRefType,
+  isRefOrRefValue,
   makeInstructionId,
 } from '../HIR';
 import {deadCodeElimination} from '../Optimization';
@@ -139,7 +138,7 @@ function infer(
       name = dep.identifier.name;
     }
 
-    if (isUseRefType(dep.identifier) || isRefValueType(dep.identifier)) {
+    if (isRefOrRefValue(dep.identifier)) {
       /*
        * TODO: this is a hack to ensure we treat functions which reference refs
        * as having a capture and therefore being considered mutable. this ensures

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableLifetimes.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableLifetimes.ts
@@ -11,6 +11,7 @@ import {
   Identifier,
   InstructionId,
   InstructionKind,
+  isRefOrRefValue,
   makeInstructionId,
   Place,
 } from '../HIR/HIR';
@@ -66,7 +67,9 @@ import {assertExhaustive} from '../Utils/utils';
  */
 
 function infer(place: Place, instrId: InstructionId): void {
-  place.identifier.mutableRange.end = makeInstructionId(instrId + 1);
+  if (!isRefOrRefValue(place.identifier)) {
+    place.identifier.mutableRange.end = makeInstructionId(instrId + 1);
+  }
 }
 
 function inferPlace(
@@ -171,7 +174,10 @@ export function inferMutableLifetimes(
         const declaration = contextVariableDeclarationInstructions.get(
           instr.value.lvalue.place.identifier,
         );
-        if (declaration != null) {
+        if (
+          declaration != null &&
+          !isRefOrRefValue(instr.value.lvalue.place.identifier)
+        ) {
           const range = instr.value.lvalue.place.identifier.mutableRange;
           if (range.start === 0) {
             range.start = declaration;

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRangesForAlias.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRangesForAlias.ts
@@ -5,7 +5,12 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {HIRFunction, Identifier, InstructionId} from '../HIR/HIR';
+import {
+  HIRFunction,
+  Identifier,
+  InstructionId,
+  isRefOrRefValue,
+} from '../HIR/HIR';
 import DisjointSet from '../Utils/DisjointSet';
 
 export function inferMutableRangesForAlias(
@@ -19,7 +24,8 @@ export function inferMutableRangesForAlias(
      * mutated.
      */
     const mutatingIdentifiers = [...aliasSet].filter(
-      id => id.mutableRange.end - id.mutableRange.start > 1,
+      id =>
+        id.mutableRange.end - id.mutableRange.start > 1 && !isRefOrRefValue(id),
     );
 
     if (mutatingIdentifiers.length > 0) {
@@ -36,7 +42,10 @@ export function inferMutableRangesForAlias(
        * last mutation.
        */
       for (const alias of aliasSet) {
-        if (alias.mutableRange.end < lastMutatingInstructionId) {
+        if (
+          alias.mutableRange.end < lastMutatingInstructionId &&
+          !isRefOrRefValue(alias)
+        ) {
           alias.mutableRange.end = lastMutatingInstructionId as InstructionId;
         }
       }

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
@@ -30,8 +30,7 @@ import {
   isArrayType,
   isMutableEffect,
   isObjectType,
-  isRefValueType,
-  isUseRefType,
+  isRefOrRefValue,
 } from '../HIR/HIR';
 import {FunctionSignature} from '../HIR/ObjectShape';
 import {
@@ -523,10 +522,7 @@ class InferenceState {
         break;
       }
       case Effect.Mutate: {
-        if (
-          isRefValueType(place.identifier) ||
-          isUseRefType(place.identifier)
-        ) {
+        if (isRefOrRefValue(place.identifier)) {
           // no-op: refs are validate via ValidateNoRefAccessInRender
         } else if (valueKind.kind === ValueKind.Context) {
           functionEffect = {
@@ -567,10 +563,7 @@ class InferenceState {
         break;
       }
       case Effect.Store: {
-        if (
-          isRefValueType(place.identifier) ||
-          isUseRefType(place.identifier)
-        ) {
+        if (isRefOrRefValue(place.identifier)) {
           // no-op: refs are validate via ValidateNoRefAccessInRender
         } else if (valueKind.kind === ValueKind.Context) {
           functionEffect = {

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts b/compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoRefAccesInRender.ts
@@ -11,6 +11,7 @@ import {
   IdentifierId,
   Place,
   SourceLocation,
+  isRefOrRefValue,
   isRefValueType,
   isUseRefType,
 } from '../HIR';
@@ -231,8 +232,7 @@ function validateNoRefAccess(
   loc: SourceLocation,
 ): void {
   if (
-    isRefValueType(operand.identifier) ||
-    isUseRefType(operand.identifier) ||
+    isRefOrRefValue(operand.identifier) ||
     refAccessingFunctions.has(operand.identifier.id)
   ) {
     errors.push({

diff --git a/...mpiler/src/__tests__/fixtures/compiler/capture-ref-for-later-mutation.expect.md b/...mpiler/src/__tests__/fixtures/compiler/capture-ref-for-later-mutation.expect.md
@@ -36,27 +36,26 @@ import { useRef } from "react";
 import { addOne } from "shared-runtime";
 
 function useKeyCommand() {
-  const $ = _c(2);
+  const $ = _c(1);
   const currentPosition = useRef(0);
-  const handleKey = (direction) => () => {
-    const position = currentPosition.current;
-    const nextPosition = direction === "left" ? addOne(position) : position;
-    currentPosition.current = nextPosition;
-  };
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const handleKey = (direction) => () => {
+      const position = currentPosition.current;
+      const nextPosition = direction === "left" ? addOne(position) : position;
+      currentPosition.current = nextPosition;
+    };
 
-  const moveLeft = { handler: handleKey("left") };
+    const moveLeft = { handler: handleKey("left") };
 
-  const t0 = handleKey("right");
-  let t1;
-  if ($[0] !== t0) {
-    t1 = { handler: t0 };
+    const moveRight = { handler: handleKey("right") };
+
+    t0 = [moveLeft, moveRight];
     $[0] = t0;
-    $[1] = t1;
   } else {
-    t1 = $[1];
+    t0 = $[0];
   }
-  const moveRight = t1;
-  return [moveLeft, moveRight];
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

diff --git a/...lback-set-ref-nested-property-ref-modified-later-preserve-memoization.expect.md b/...lback-set-ref-nested-property-ref-modified-later-preserve-memoization.expect.md
@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enablePreserveExistingMemoizationGuarantees
+// @enablePreserveExistingMemoizationGuarantees @validateRefAccessDuringRender
 import {useCallback, useRef} from 'react';
 
 function Component(props) {
@@ -42,7 +42,9 @@ export const FIXTURE_ENTRYPOINT = {
 > 10 |     ref.current.inner = event.target.value;
      | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 > 11 |   });
-     | ^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (7:11)
+     | ^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef). Cannot access ref value at freeze $44:TObject<BuiltInFunction> (7:11)
+
+InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (14:14)
   12 |
   13 |   // The ref is modified later, extending its range and preventing memoization of onChange
   14 |   ref.current.inner = null;

diff --git a/...error.todo-useCallback-set-ref-nested-property-ref-modified-later-preserve-memoization.js b/...error.todo-useCallback-set-ref-nested-property-ref-modified-later-preserve-memoization.js
@@ -1,4 +1,4 @@
-// @enablePreserveExistingMemoizationGuarantees
+// @enablePreserveExistingMemoizationGuarantees @validateRefAccessDuringRender
 import {useCallback, useRef} from 'react';
 
 function Component(props) {

diff --git a/...a-function-preserve-memoization.expect.md → ...a-function-preserve-memoization.expect.md b/...a-function-preserve-memoization.expect.md → ...a-function-preserve-memoization.expect.md
@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enablePreserveExistingMemoizationGuarantees
+// @enablePreserveExistingMemoizationGuarantees @validateRefAccessDuringRender
 import {useCallback, useRef} from 'react';
 
 function Component(props) {
@@ -45,7 +45,11 @@ export const FIXTURE_ENTRYPOINT = {
 > 10 |     ref.current.inner = event.target.value;
      | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 > 11 |   });
-     | ^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (7:11)
+     | ^^^^ InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef). Cannot access ref value at freeze $53:TObject<BuiltInFunction> (7:11)
+
+InvalidReact: This function accesses a ref value (the `current` property), which may not be accessed during render. (https://react.dev/reference/react/useRef). Function mutate? $77[20:22]:TObject<BuiltInFunction> accesses a ref (17:17)
+
+InvalidReact: Ref values (the `current` property) may not be accessed during render. (https://react.dev/reference/react/useRef) (17:17)
   12 |
   13 |   // The ref is modified later, extending its range and preventing memoization of onChange
   14 |   const reset = () => {

diff --git a/...ater-via-function-preserve-memoization.js → ...ater-via-function-preserve-memoization.js b/...ater-via-function-preserve-memoization.js → ...ater-via-function-preserve-memoization.js
@@ -1,4 +1,4 @@
-// @enablePreserveExistingMemoizationGuarantees
+// @enablePreserveExistingMemoizationGuarantees @validateRefAccessDuringRender
 import {useCallback, useRef} from 'react';
 
 function Component(props) {

diff --git a/...compiler/original-reactive-scopes-fork/capture-ref-for-later-mutation.expect.md b/...compiler/original-reactive-scopes-fork/capture-ref-for-later-mutation.expect.md
@@ -37,42 +37,26 @@ import { useRef } from "react";
 import { addOne } from "shared-runtime";
 
 function useKeyCommand() {
-  const $ = _c(6);
+  const $ = _c(1);
   const currentPosition = useRef(0);
-  const handleKey = (direction) => () => {
-    const position = currentPosition.current;
-    const nextPosition = direction === "left" ? addOne(position) : position;
-    currentPosition.current = nextPosition;
-  };
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = { handler: handleKey("left") };
+    const handleKey = (direction) => () => {
+      const position = currentPosition.current;
+      const nextPosition = direction === "left" ? addOne(position) : position;
+      currentPosition.current = nextPosition;
+    };
+
+    const moveLeft = { handler: handleKey("left") };
+
+    const moveRight = { handler: handleKey("right") };
+
+    t0 = [moveLeft, moveRight];
     $[0] = t0;
   } else {
     t0 = $[0];
   }
-  const moveLeft = t0;
-
-  const t1 = handleKey("right");
-  let t2;
-  if ($[1] !== t1) {
-    t2 = { handler: t1 };
-    $[1] = t1;
-    $[2] = t2;
-  } else {
-    t2 = $[2];
-  }
-  const moveRight = t2;
-  let t3;
-  if ($[3] !== moveLeft || $[4] !== moveRight) {
-    t3 = [moveLeft, moveRight];
-    $[3] = moveLeft;
-    $[4] = moveRight;
-    $[5] = t3;
-  } else {
-    t3 = $[5];
-  }
-  return t3;
+  return t0;
 }
 
 export const FIXTURE_ENTRYPOINT = {

diff --git a/...mpiler/preserve-memo-validation/error.maybe-mutable-ref-not-preserved.expect.md b/...mpiler/preserve-memo-validation/error.maybe-mutable-ref-not-preserved.expect.md
diff --git a/...res/compiler/preserve-memo-validation/maybe-mutable-ref-not-preserved.expect.md b/...res/compiler/preserve-memo-validation/maybe-mutable-ref-not-preserved.expect.md
@@ -0,0 +1,53 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees:true
+
+import {useRef, useMemo} from 'react';
+import {makeArray} from 'shared-runtime';
+
+function useFoo() {
+  const r = useRef();
+  return useMemo(() => makeArray(r), []);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validatePreserveExistingMemoizationGuarantees:true
+
+import { useRef, useMemo } from "react";
+import { makeArray } from "shared-runtime";
+
+function useFoo() {
+  const $ = _c(1);
+  const r = useRef();
+  let t0;
+  let t1;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = makeArray(r);
+    $[0] = t1;
+  } else {
+    t1 = $[0];
+  }
+  t0 = t1;
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) [{}]
diff --git a/.../error.maybe-mutable-ref-not-preserved.ts → ...dation/maybe-mutable-ref-not-preserved.ts b/.../error.maybe-mutable-ref-not-preserved.ts → ...dation/maybe-mutable-ref-not-preserved.ts