[compiler] fix false positive "mutate frozen" validation with refs

The test case here previously reported a "Cannot modify local variables after render completes" error (from ValidateNoFreezingKnownMutableFunctions). This happens because one of the functions passed to a hook clearly mutates a ref — except that we try to ignore mutations of refs! The problem in this case is that the `const ref = ...` was getting converted to a context variable since the ref is accessed in a function before its declaration. We don't infer types for context variables at all, and our ref handling is based on types, so we failed to ignore this ref mutation.

The fix is to recognize that `StoreLocal const ...` is a special case: the variable may be referenced in code before the declaration, but at runtime it's either a TDZ error or the variable will have the type from the declaration. So we can safely infer a type.

Author: josephsavona

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

@@ -249,6 +249,7 @@ export function inferMutationAliasingRanges(
   }
   for (const param of [...fn.context, ...fn.params]) {
     const place = param.kind === 'Identifier' ? param : param.place;
+
     const node = state.nodes.get(place.identifier);
     if (node == null) {
       continue;

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

@@ -14,6 +14,7 @@ import {
   Identifier,
   IdentifierId,
   Instruction,
+  InstructionKind,
   makePropertyLiteral,
   makeType,
   PropType,
@@ -194,12 +195,29 @@ function* generateInstructionTypes(
       break;
     }
 
-    // We intentionally do not infer types for context variables
+    // We intentionally do not infer types for most context variables
     case 'DeclareContext':
-    case 'StoreContext':
     case 'LoadContext': {
       break;
     }
+    case 'StoreContext': {
+      /**
+       * The caveat is StoreContext const, where we know the value is
+       * assigned once such that everywhere the value is accessed, it
+       * must have the same type from the rvalue.
+       *
+       * A concrete example where this is useful is `const ref = useRef()`
+       * where the ref is referenced before its declaration in a function
+       * expression, causing it to be converted to a const context variable.
+       */
+      if (value.lvalue.kind === InstructionKind.Const) {
+        yield equation(
+          value.lvalue.place.identifier.type,
+          value.value.identifier.type,
+        );
+      }
+      break;
+    }
 
     case 'StoreLocal': {
       if (env.config.enableUseTypeAnnotations) {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-mutate-ref-in-function-passed-to-hook.expect.md

@@ -0,0 +1,108 @@
+
+## Input
+
+```javascript
+// @flow
+component Example() {
+  const fooRef = useRef();
+
+  function updateStyles() {
+    const foo = fooRef.current;
+    if (barRef.current == null || foo == null) {
+      return;
+    }
+    foo.style.height = '100px';
+  }
+
+  const barRef = useRef(null);
+
+  const resizeRef = useResizeObserver(rect => {
+    const {width} = rect;
+    barRef.current = width;
+  });
+
+  useLayoutEffect(() => {
+    const observer = new ResizeObserver(_ => {
+      updateStyles();
+    });
+
+    return () => {
+      observer.disconnect();
+    };
+  }, []);
+
+  return <div ref={resizeRef} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Example() {
+  const $ = _c(6);
+  const fooRef = useRef();
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = function updateStyles() {
+      const foo = fooRef.current;
+      if (barRef.current == null || foo == null) {
+        return;
+      }
+
+      foo.style.height = "100px";
+    };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const updateStyles = t0;
+
+  const barRef = useRef(null);
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = (rect) => {
+      const { width } = rect;
+      barRef.current = width;
+    };
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const resizeRef = useResizeObserver(t1);
+  let t2;
+  let t3;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = () => {
+      const observer = new ResizeObserver((_) => {
+        updateStyles();
+      });
+      return () => {
+        observer.disconnect();
+      };
+    };
+
+    t3 = [];
+    $[2] = t2;
+    $[3] = t3;
+  } else {
+    t2 = $[2];
+    t3 = $[3];
+  }
+  useLayoutEffect(t2, t3);
+  let t4;
+  if ($[4] !== resizeRef) {
+    t4 = <div ref={resizeRef} />;
+    $[4] = resizeRef;
+    $[5] = t4;
+  } else {
+    t4 = $[5];
+  }
+  return t4;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-mutate-ref-in-function-passed-to-hook.js

@@ -0,0 +1,31 @@
+// @flow
+component Example() {
+  const fooRef = useRef();
+
+  function updateStyles() {
+    const foo = fooRef.current;
+    if (barRef.current == null || foo == null) {
+      return;
+    }
+    foo.style.height = '100px';
+  }
+
+  const barRef = useRef(null);
+
+  const resizeRef = useResizeObserver(rect => {
+    const {width} = rect;
+    barRef.current = width;
+  });
+
+  useLayoutEffect(() => {
+    const observer = new ResizeObserver(_ => {
+      updateStyles();
+    });
+
+    return () => {
+      observer.disconnect();
+    };
+  }, []);
+
+  return <div ref={resizeRef} />;
+}