[compiler] Validate against locals being reassigned after render

Adds a pass which validates that local variables are not reassigned by functions which may be called after render. This is a straightforward forward data-flow analysis, where we:
1. Build up a mapping of context variables in the outer component/hook
2. Find ObjectMethod/FunctionExpressions which may reassign those context variables
3. Propagate aliases of those functions via StoreLocal/LoadLocal
4. Disallow passing those functions with a Freeze effect. This includes JSX arguments, hook arguments, hook return types, etc.

Conceptually, a function that reassigns a local is inherently mutable. Frozen functions must be side-effect free, so these two categories are incompatible and we can use the freeze effect to find all instances of where such functions are disallowed rather than special-casing eg hook calls and JSX.

ghstack-source-id: c2b22e3d62
Pull Request resolved: https://github.com/facebook/react/pull/30107

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -96,6 +96,7 @@ import {
   validatePreservedManualMemoization,
   validateUseMemo,
 } from "../Validation";
+import { validateLocalsNotReassignedAfterRender } from "../Validation/ValidateLocalsNotReassignedAfterRender";
 
 export type CompilerPipelineValue =
   | { kind: "ast"; name: string; value: CodegenFunction }
@@ -202,6 +203,8 @@ function* runWithEnvironment(
   inferReferenceEffects(hir);
   yield log({ kind: "hir", name: "InferReferenceEffects", value: hir });
 
+  validateLocalsNotReassignedAfterRender(hir);
+
   // Note: Has to come after infer reference effects because "dead" code may still affect inference
   deadCodeElimination(hir);
   yield log({ kind: "hir", name: "DeadCodeElimination", value: hir });

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

@@ -0,0 +1,146 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { CompilerError, Effect } from "..";
+import { HIRFunction, IdentifierId, Place } from "../HIR";
+import {
+  eachInstructionValueOperand,
+  eachTerminalOperand,
+} from "../HIR/visitors";
+
+/**
+ * Validates that local variables cannot be reassigned after render.
+ * This prevents a category of bugs in which a closure captures a
+ * binding from one render but does not update
+ */
+export function validateLocalsNotReassignedAfterRender(fn: HIRFunction): void {
+  const contextVariables = new Set<IdentifierId>();
+  const reassignment = getContextReassignment(fn, contextVariables, false);
+  if (reassignment !== null) {
+    CompilerError.throwInvalidReact({
+      reason:
+        "Reassigning a variable after render has completed can cause inconsistent behavior on subsequent renders. Consider using state instead",
+      description:
+        reassignment.identifier.name !== null &&
+        reassignment.identifier.name.kind === "named"
+          ? `Variable \`${reassignment.identifier.name.value}\` cannot be reassigned after render`
+          : "",
+      loc: reassignment.loc,
+    });
+  }
+}
+
+function getContextReassignment(
+  fn: HIRFunction,
+  contextVariables: Set<IdentifierId>,
+  isFunctionExpression: boolean
+): Place | null {
+  const reassigningFunctions = new Map<IdentifierId, Place>();
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      const { lvalue, value } = instr;
+      switch (value.kind) {
+        case "FunctionExpression":
+        case "ObjectMethod": {
+          let reassignment = getContextReassignment(
+            value.loweredFunc.func,
+            contextVariables,
+            true
+          );
+          if (reassignment === null) {
+            // If the function itself doesn't reassign, does one of its dependencies?
+            for (const operand of eachInstructionValueOperand(value)) {
+              const reassignmentFromOperand = reassigningFunctions.get(
+                operand.identifier.id
+              );
+              if (reassignmentFromOperand !== undefined) {
+                reassignment = reassignmentFromOperand;
+                break;
+              }
+            }
+          }
+          // if the function or its depends reassign, propagate that fact on the lvalue
+          if (reassignment !== null) {
+            reassigningFunctions.set(lvalue.identifier.id, reassignment);
+          }
+          break;
+        }
+        case "StoreLocal": {
+          const reassignment = reassigningFunctions.get(
+            value.value.identifier.id
+          );
+          if (reassignment !== undefined) {
+            reassigningFunctions.set(
+              value.lvalue.place.identifier.id,
+              reassignment
+            );
+            reassigningFunctions.set(lvalue.identifier.id, reassignment);
+          }
+          break;
+        }
+        case "LoadLocal": {
+          const reassignment = reassigningFunctions.get(
+            value.place.identifier.id
+          );
+          if (reassignment !== undefined) {
+            reassigningFunctions.set(lvalue.identifier.id, reassignment);
+          }
+          break;
+        }
+        case "DeclareContext": {
+          if (!isFunctionExpression) {
+            contextVariables.add(value.lvalue.place.identifier.id);
+          }
+          break;
+        }
+        case "StoreContext": {
+          if (isFunctionExpression) {
+            if (contextVariables.has(value.lvalue.place.identifier.id)) {
+              return value.lvalue.place;
+            }
+          } else {
+            /*
+             * We only track reassignments of variables defined in the outer
+             * component or hook.
+             */
+            contextVariables.add(value.lvalue.place.identifier.id);
+          }
+          break;
+        }
+        default: {
+          for (const operand of eachInstructionValueOperand(value)) {
+            CompilerError.invariant(operand.effect !== Effect.Unknown, {
+              reason: `Expected effects to be inferred prior to ValidateLocalsNotReassignedAfterRender`,
+              loc: operand.loc,
+            });
+            const reassignment = reassigningFunctions.get(
+              operand.identifier.id
+            );
+            if (
+              reassignment !== undefined &&
+              operand.effect === Effect.Freeze
+            ) {
+              /*
+               * Functions that reassign local variables are inherently mutable and are unsafe to pass
+               * to a place that expects a frozen value. Propagate the reassignment upward.
+               */
+              return reassignment;
+            }
+          }
+          break;
+        }
+      }
+    }
+    for (const operand of eachTerminalOperand(block.terminal)) {
+      const reassignment = reassigningFunctions.get(operand.identifier.id);
+      if (reassignment !== undefined) {
+        return reassignment;
+      }
+    }
+  }
+  return null;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-reassign-local-in-hook-return-value.expect.md

@@ -0,0 +1,27 @@
+
+## Input
+
+```javascript
+function useFoo() {
+  let x = 0;
+  return (value) => {
+    x = value;
+  };
+}
+
+```
+
+
+## Error
+
+```
+  2 |   let x = 0;
+  3 |   return (value) => {
+> 4 |     x = value;
+    |     ^ InvalidReact: Reassigning a variable after render has completed can cause inconsistent behavior on subsequent renders. Consider using state instead. Variable `x` cannot be reassigned after render (4:4)
+  5 |   };
+  6 | }
+  7 |
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-reassign-local-in-hook-return-value.js

@@ -0,0 +1,6 @@
+function useFoo() {
+  let x = 0;
+  return (value) => {
+    x = value;
+  };
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-reassign-local-variable-in-effect.expect.md

@@ -43,56 +43,17 @@ function Component() {
 
 ```
 
-## Code
 
-```javascript
-import { c as _c } from "react/compiler-runtime";
-import { useEffect } from "react";
-
-function Component() {
-  const $ = _c(4);
-  let local;
-  let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = (newValue) => {
-      local = newValue;
-    };
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  const reassignLocal = t0;
-  let t1;
-  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
-    t1 = (newValue_0) => {
-      reassignLocal("hello");
-      if (local === newValue_0) {
-        console.log("`local` was updated!");
-      } else {
-        throw new Error("`local` not updated!");
-      }
-    };
-    $[1] = t1;
-  } else {
-    t1 = $[1];
-  }
-  const onMount = t1;
-  let t2;
-  let t3;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
-    t2 = () => {
-      onMount();
-    };
-    t3 = [onMount];
-    $[2] = t2;
-    $[3] = t3;
-  } else {
-    t2 = $[2];
-    t3 = $[3];
-  }
-  useEffect(t2, t3);
-  return "ok";
-}
+## Error
 
 ```
+   5 |
+   6 |   const reassignLocal = (newValue) => {
+>  7 |     local = newValue;
+     |     ^^^^^ InvalidReact: Reassigning a variable after render has completed can cause inconsistent behavior on subsequent renders. Consider using state instead. Variable `local` cannot be reassigned after render (7:7)
+   8 |   };
+   9 |
+  10 |   const onMount = (newValue) => {
+```
+          
       

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-reassign-local-variable-in-hook-argument.expect.md

@@ -44,53 +44,17 @@ function Component() {
 
 ```
 
-## Code
 
-```javascript
-import { c as _c } from "react/compiler-runtime";
-import { useEffect } from "react";
-import { useIdentity } from "shared-runtime";
-
-function Component() {
-  const $ = _c(3);
-  let local;
-  let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = (newValue) => {
-      local = newValue;
-    };
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  const reassignLocal = t0;
-  let t1;
-  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
-    t1 = (newValue_0) => {
-      reassignLocal("hello");
-      if (local === newValue_0) {
-        console.log("`local` was updated!");
-      } else {
-        throw new Error("`local` not updated!");
-      }
-    };
-    $[1] = t1;
-  } else {
-    t1 = $[1];
-  }
-  const callback = t1;
-  let t2;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
-    t2 = () => {
-      callback();
-    };
-    $[2] = t2;
-  } else {
-    t2 = $[2];
-  }
-  useIdentity(t2);
-  return "ok";
-}
+## Error
 
 ```
+   6 |
+   7 |   const reassignLocal = (newValue) => {
+>  8 |     local = newValue;
+     |     ^^^^^ InvalidReact: Reassigning a variable after render has completed can cause inconsistent behavior on subsequent renders. Consider using state instead. Variable `local` cannot be reassigned after render (8:8)
+   9 |   };
+  10 |
+  11 |   const callback = (newValue) => {
+```
+          
       

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-reassign-local-variable-in-jsx-callback.expect.md

@@ -37,41 +37,17 @@ function Component() {
 
 ```
 
-## Code
 
-```javascript
-import { c as _c } from "react/compiler-runtime";
-function Component() {
-  const $ = _c(2);
-  let local;
-  let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = (newValue) => {
-      local = newValue;
-    };
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  const reassignLocal = t0;
-  let t1;
-  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
-    const onClick = (newValue_0) => {
-      reassignLocal("hello");
-      if (local === newValue_0) {
-        console.log("`local` was updated!");
-      } else {
-        throw new Error("`local` not updated!");
-      }
-    };
-
-    t1 = <button onClick={onClick}>Submit</button>;
-    $[1] = t1;
-  } else {
-    t1 = $[1];
-  }
-  return t1;
-}
+## Error
 
 ```
+  3 |
+  4 |   const reassignLocal = (newValue) => {
+> 5 |     local = newValue;
+    |     ^^^^^ InvalidReact: Reassigning a variable after render has completed can cause inconsistent behavior on subsequent renders. Consider using state instead. Variable `local` cannot be reassigned after render (5:5)
+  6 |   };
+  7 |
+  8 |   const onClick = (newValue) => {
+```
+          
       

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-mutable-range-shared-inner-outer-function.expect.md

@@ -18,7 +18,7 @@ function Component(props) {
     a.property = true;
     b.push(false);
   };
-  return <div onClick={f} />;
+  return <div onClick={f()} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -35,10 +35,10 @@ import { c as _c } from "react/compiler-runtime"; // @enableAssumeHooksFollowRul
 let cond = true;
 function Component(props) {
   const $ = _c(1);
-  let a;
-  let b;
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    let a;
+    let b;
     const f = () => {
       if (cond) {
         a = {};
@@ -52,7 +52,7 @@ function Component(props) {
       b.push(false);
     };
 
-    t0 = <div onClick={f} />;
+    t0 = <div onClick={f()} />;
     $[0] = t0;
   } else {
     t0 = $[0];

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-mutable-range-shared-inner-outer-function.js

@@ -14,7 +14,7 @@ function Component(props) {
     a.property = true;
     b.push(false);
   };
-  return <div onClick={f} />;
+  return <div onClick={f()} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {