[compiler] fix false positive "mutate frozen" validation with refs (facebook#33993)

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.

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

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,114 @@
+
+## Input
+
+```javascript
+// @flow
+component Example() {
+  const fooRef = useRef();
+
+  function updateStyles() {
+    const foo = fooRef.current;
+    // The access of `barRef` here before its declaration causes it be hoisted...
+    if (barRef.current == null || foo == null) {
+      return;
+    }
+    foo.style.height = '100px';
+  }
+
+  // ...which previously meant that we didn't infer a type...
+  const barRef = useRef(null);
+
+  const resizeRef = useResizeObserver(
+    rect => {
+      const {width} = rect;
+      // ...which meant that we failed to ignore the mutation here...
+      barRef.current = width;
+    } // ...which caused this to fail with "can't freeze a mutable function"
+  );
+
+  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,36 @@
+// @flow
+component Example() {
+  const fooRef = useRef();
+
+  function updateStyles() {
+    const foo = fooRef.current;
+    // The access of `barRef` here before its declaration causes it be hoisted...
+    if (barRef.current == null || foo == null) {
+      return;
+    }
+    foo.style.height = '100px';
+  }
+
+  // ...which previously meant that we didn't infer a type...
+  const barRef = useRef(null);
+
+  const resizeRef = useResizeObserver(
+    rect => {
+      const {width} = rect;
+      // ...which meant that we failed to ignore the mutation here...
+      barRef.current = width;
+    } // ...which caused this to fail with "can't freeze a mutable function"
+  );
+
+  useLayoutEffect(() => {
+    const observer = new ResizeObserver(_ => {
+      updateStyles();
+    });
+
+    return () => {
+      observer.disconnect();
+    };
+  }, []);
+
+  return <div ref={resizeRef} />;
+}