[compiler] Patch for reactive refs in inferred effect dependencies

compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts

@@ -1738,6 +1738,40 @@ export function isStableType(id: Identifier): boolean {
   );
 }
 
+export function isStableTypeContainer(id: Identifier): boolean {
+  const type_ = id.type;
+  if (type_.kind !== 'Object') {
+    return false;
+  }
+  return (
+    isUseStateType(id) || // setState
+    type_.shapeId === 'BuiltInUseActionState' || // setActionState
+    isUseReducerType(id) || // dispatcher
+    type_.shapeId === 'BuiltInUseTransition' // startTransition
+  );
+}
+
+export function evaluatesToStableTypeOrContainer(
+  env: Environment,
+  {value}: Instruction,
+): boolean {
+  if (value.kind === 'CallExpression' || value.kind === 'MethodCall') {
+    const callee =
+      value.kind === 'CallExpression' ? value.callee : value.property;
+
+    const calleeHookKind = getHookKind(env, callee.identifier);
+    switch (calleeHookKind) {
+      case 'useState':
+      case 'useReducer':
+      case 'useActionState':
+      case 'useRef':
+      case 'useTransition':
+        return true;
+    }
+  }
+  return false;
+}
+
 export function isUseEffectHookType(id: Identifier): boolean {
   return (
     id.type.kind === 'Function' && id.type.shapeId === 'BuiltInUseEffectHook'

compiler/packages/babel-plugin-react-compiler/src/Inference/InferReactivePlaces.ts

@@ -9,14 +9,19 @@ import {CompilerError} from '..';
 import {
   BlockId,
   Effect,
+  Environment,
   HIRFunction,
   Identifier,
   IdentifierId,
+  Instruction,
   Place,
   computePostDominatorTree,
+  evaluatesToStableTypeOrContainer,
   getHookKind,
   isStableType,
+  isStableTypeContainer,
   isUseOperator,
+  isUseRefType,
 } from '../HIR';
 import {PostDominator} from '../HIR/Dominator';
 import {
@@ -31,6 +36,103 @@ import {
 import DisjointSet from '../Utils/DisjointSet';
 import {assertExhaustive} from '../Utils/utils';
 
+/**
+ * Side map to track and propagate sources of stability (i.e. hook calls such as
+ * `useRef()` and property reads such as `useState()[1]). Note that this
+ * requires forward data flow analysis since stability is not part of React
+ * Compiler's type system.
+ */
+class StableSidemap {
+  map: Map<IdentifierId, {isStable: boolean}> = new Map();
+  env: Environment;
+
+  constructor(env: Environment) {
+    this.env = env;
+  }
+
+  handleInstruction(instr: Instruction): void {
+    const {value, lvalue} = instr;
+
+    switch (value.kind) {
+      case 'CallExpression':
+      case 'MethodCall': {
+        /**
+         * Sources of stability are known hook calls
+         */
+        if (evaluatesToStableTypeOrContainer(this.env, instr)) {
+          if (isStableType(lvalue.identifier)) {
+            this.map.set(lvalue.identifier.id, {
+              isStable: true,
+            });
+          } else {
+            this.map.set(lvalue.identifier.id, {
+              isStable: false,
+            });
+          }
+        } else if (
+          this.env.config.enableTreatRefLikeIdentifiersAsRefs &&
+          isUseRefType(lvalue.identifier)
+        ) {
+          this.map.set(lvalue.identifier.id, {
+            isStable: true,
+          });
+        }
+        break;
+      }
+
+      case 'Destructure':
+      case 'PropertyLoad': {
+        /**
+         * PropertyLoads may from stable containers may also produce stable
+         * values. ComputedLoads are technically safe for now (as all stable
+         * containers have differently-typed elements), but are not handled as
+         * they should be rare anyways.
+         */
+        const source =
+          value.kind === 'Destructure'
+            ? value.value.identifier.id
+            : value.object.identifier.id;
+        const entry = this.map.get(source);
+        if (entry) {
+          for (const lvalue of eachInstructionLValue(instr)) {
+            if (isStableTypeContainer(lvalue.identifier)) {
+              this.map.set(lvalue.identifier.id, {
+                isStable: false,
+              });
+            } else if (isStableType(lvalue.identifier)) {
+              this.map.set(lvalue.identifier.id, {
+                isStable: true,
+              });
+            }
+          }
+        }
+        break;
+      }
+
+      case 'StoreLocal': {
+        const entry = this.map.get(value.value.identifier.id);
+        if (entry) {
+          this.map.set(lvalue.identifier.id, entry);
+          this.map.set(value.lvalue.place.identifier.id, entry);
+        }
+        break;
+      }
+
+      case 'LoadLocal': {
+        const entry = this.map.get(value.place.identifier.id);
+        if (entry) {
+          this.map.set(lvalue.identifier.id, entry);
+        }
+        break;
+      }
+    }
+  }
+
+  isStable(id: IdentifierId): boolean {
+    const entry = this.map.get(id);
+    return entry != null ? entry.isStable : false;
+  }
+}
 /*
  * Infers which `Place`s are reactive, ie may *semantically* change
  * over the course of the component/hook's lifetime. Places are reactive
@@ -111,6 +213,7 @@ import {assertExhaustive} from '../Utils/utils';
  */
 export function inferReactivePlaces(fn: HIRFunction): void {
   const reactiveIdentifiers = new ReactivityMap(findDisjointMutableValues(fn));
+  const stableIdentifierSources = new StableSidemap(fn.env);
   for (const param of fn.params) {
     const place = param.kind === 'Identifier' ? param : param.place;
     reactiveIdentifiers.markReactive(place);
@@ -184,6 +287,7 @@ export function inferReactivePlaces(fn: HIRFunction): void {
         }
       }
       for (const instruction of block.instructions) {
+        stableIdentifierSources.handleInstruction(instruction);
         const {value} = instruction;
         let hasReactiveInput = false;
         /*
@@ -218,7 +322,13 @@ export function inferReactivePlaces(fn: HIRFunction): void {
 
         if (hasReactiveInput) {
           for (const lvalue of eachInstructionLValue(instruction)) {
-            if (isStableType(lvalue.identifier)) {
+            /**
+             * Note that it's not correct to mark all stable-typed identifiers
+             * as non-reactive, since ternaries and other value blocks can
+             * produce reactive identifiers typed as these.
+             * (e.g. `props.cond ? setState1 : setState2`)
+             */
+            if (stableIdentifierSources.isStable(lvalue.identifier.id)) {
               continue;
             }
             reactiveIdentifiers.markReactive(lvalue);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-ref-ternary.expect.md

@@ -0,0 +1,69 @@
+
+## Input
+
+```javascript
+// @inferEffectDependencies
+import {useRef, useEffect} from 'react';
+import {print, mutate} from 'shared-runtime';
+
+function Component({cond}) {
+  const arr = useRef([]);
+  const other = useRef([]);
+  // Although arr and other are both stable, derived is not
+  const derived = cond ? arr : other;
+  useEffect(() => {
+    mutate(derived.current);
+    print(derived.current);
+  });
+  return arr;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @inferEffectDependencies
+import { useRef, useEffect } from "react";
+import { print, mutate } from "shared-runtime";
+
+function Component(t0) {
+  const $ = _c(4);
+  const { cond } = t0;
+  let t1;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = [];
+    $[0] = t1;
+  } else {
+    t1 = $[0];
+  }
+  const arr = useRef(t1);
+  let t2;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = [];
+    $[1] = t2;
+  } else {
+    t2 = $[1];
+  }
+  const other = useRef(t2);
+
+  const derived = cond ? arr : other;
+  let t3;
+  if ($[2] !== derived) {
+    t3 = () => {
+      mutate(derived.current);
+      print(derived.current);
+    };
+    $[2] = derived;
+    $[3] = t3;
+  } else {
+    t3 = $[3];
+  }
+  useEffect(t3, [derived]);
+  return arr;
+}
+
+```
+
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-ref-ternary.js

@@ -0,0 +1,15 @@
+// @inferEffectDependencies
+import {useRef, useEffect} from 'react';
+import {print, mutate} from 'shared-runtime';
+
+function Component({cond}) {
+  const arr = useRef([]);
+  const other = useRef([]);
+  // Although arr and other are both stable, derived is not
+  const derived = cond ? arr : other;
+  useEffect(() => {
+    mutate(derived.current);
+    print(derived.current);
+  });
+  return arr;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md

@@ -83,10 +83,10 @@ export const FIXTURE_ENTRYPOINT = {
 import { c as _c2 } from "react/compiler-runtime"; // @inlineJsxTransform
 
 function Parent(t0) {
-  const $ = _c2(2);
+  const $ = _c2(3);
   const { children, ref } = t0;
   let t1;
-  if ($[0] !== children) {
+  if ($[0] !== children || $[1] !== ref) {
     if (DEV) {
       t1 = <div ref={ref}>{children}</div>;
     } else {
@@ -99,9 +99,10 @@ function Parent(t0) {
       };
     }
     $[0] = children;
-    $[1] = t1;
+    $[1] = ref;
+    $[2] = t1;
   } else {
-    t1 = $[1];
+    t1 = $[2];
   }
   return t1;
 }