Validate frozen lambdas are actually frozen

Adds validation to reject freezing mutable lambdas, since lambdas cannot _be_ 
frozen, they're either frozen or not. Example invalid code: 

```javascript 

function Component(props) { 

const x = {value: ""}; 

const onChange = (e) => { 

// MUTATION!!!!! 

x.value = e.target.value; 

setX(x); 

}; 

return <input value={x.value} onChange={onChange} />; 

} 

``` 

Note that there is a separate issue in which we are not detecting lambdas that 
would definitely modify immutable values. We may need to distinguish 
ConditionallyCapture (captures if the value is mutable, otherwise readonly) from 
Capture (definitely mutates) in order to make that case work. But already this 
validation helps prevent some invalid code.

compiler/forget/packages/snap/src/compiler-worker.ts

@@ -159,6 +159,7 @@ export async function compile(
         memoizeJsxElements,
         validateHooksUsage: true,
         validateRefAccessDuringRender,
+        validateFrozenLambdas: true,
       },
       logger: null,
       gating,

compiler/forget/src/CompilerPipeline.ts

@@ -13,6 +13,7 @@ import {
   mergeConsecutiveBlocks,
   ReactiveFunction,
   validateConsistentIdentifiers,
+  validateFrozenLambdas,
   validateHooksUsage,
   validateNoRefAccessInRender,
   validateTerminalSuccessors,
@@ -116,6 +117,10 @@ export function* run(
   inferReferenceEffects(hir);
   yield log({ kind: "hir", name: "InferReferenceEffects", value: hir });
 
+  if (env.validateFrozenLambdas) {
+    validateFrozenLambdas(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/forget/src/HIR/Environment.ts

@@ -74,6 +74,14 @@ export type EnvironmentConfig = Partial<{
    */
   validateRefAccessDuringRender: boolean;
 
+  /**
+   * Validate that mutable lambdas are not passed where a frozen value is expected, since mutable
+   * lambdas cannot be frozen. The only mutation allowed inside a frozen lambda is of ref values.
+   *
+   * Defaults to false
+   */
+  validateFrozenLambdas: boolean;
+
   /**
    * Enable inlining of `useMemo()` function expressions so that they can be more optimally
    * compiled.
@@ -127,6 +135,7 @@ export class Environment {
   #nextBlock: number = 0;
   validateHooksUsage: boolean;
   validateRefAccessDuringRender: boolean;
+  validateFrozenLambdas: boolean;
   enableFunctionCallSignatureOptimizations: boolean;
   enableAssumeHooksFollowRulesOfReact: boolean;
   enableTreatHooksAsFunctions: boolean;
@@ -164,6 +173,7 @@ export class Environment {
     this.validateHooksUsage = config?.validateHooksUsage ?? false;
     this.validateRefAccessDuringRender =
       config?.validateRefAccessDuringRender ?? false;
+    this.validateFrozenLambdas = config?.validateFrozenLambdas ?? false;
     this.enableFunctionCallSignatureOptimizations =
       config?.enableFunctionCallSignatureOptimizations ?? false;
     this.enableAssumeHooksFollowRulesOfReact =

compiler/forget/src/HIR/ValidateFrozenLambdas.ts

@@ -0,0 +1,102 @@
+/**
+ * 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,
+  CompilerErrorDetail,
+  ErrorSeverity,
+} from "../CompilerError";
+import {
+  Effect,
+  FunctionExpression,
+  HIRFunction,
+  IdentifierId,
+  isMutableEffect,
+  isRefValueType,
+  isUseRefType,
+} from "./HIR";
+import { eachInstructionValueOperand } from "./visitors";
+
+/**
+ * Various APIs in React take ownership of the values passed to them, such that it is invalid
+ * to subsequently modify those values. Examples include:
+ * - Passing a value as a prop to JSX. Subsequently mutating this value will result in undefined
+ *   behavior, since the mutation may or may not be observed depending on when the child re-renders.
+ *   In addition, the value may be used as an input to memoization in children, and mutation could
+ *   invalidate that memoization.
+ * - Passing a value to `useState()`, for the same reason.
+ * - Passing a value to a hook, for the same reason.
+ *
+ * Most "normal" data types (objects, arrays, etc) can be "frozen" when passed to a React API simply
+ * by not calling any mutating methods on them. However, mutable lambdas are an exception: if a lambda
+ * has side-effects, there is no way to do something to the lambda that would allow calling it without
+ * triggering those side effects. The only thing a developer could do is not call the lambda, but developers
+ * also have no way of knowing that they can't call the lambda.
+ *
+ * From a type system perspective, the above APIs that "take ownership" of their values really accept
+ * *already frozen* values as input. Thus it is invalid to pass a value that cannot be frozen to these APIs,
+ * and it is therefore invalid to pass a mutable lambda.
+ *
+ * This pass validates the above rule. Note that this validation can by bypassed by storing a mutable lambda
+ * inside some value (eg as an array element or object property). In these cases we trust that the developer
+ * is not breaking the rules. The goal of this validation is to find cases that are provably wrong and help
+ * the developer fix the mistake earlier.
+ */
+export function validateFrozenLambdas(fn: HIRFunction): void {
+  const lambdas = new Map<IdentifierId, FunctionExpression>();
+  const temporaries = new Map<IdentifierId, IdentifierId>();
+
+  const errors = new CompilerError();
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      if (instr.value.kind === "FunctionExpression") {
+        lambdas.set(instr.lvalue.identifier.id, instr.value);
+      } else if (instr.value.kind === "LoadLocal") {
+        const resolvedId =
+          temporaries.get(instr.value.place.identifier.id) ??
+          instr.value.place.identifier.id;
+        temporaries.set(instr.lvalue.identifier.id, resolvedId);
+      } else if (instr.value.kind === "StoreLocal") {
+        const resolvedId =
+          temporaries.get(instr.value.value.identifier.id) ??
+          instr.value.value.identifier.id;
+        temporaries.set(instr.value.lvalue.place.identifier.id, resolvedId);
+      } else {
+        for (const operand of eachInstructionValueOperand(instr.value)) {
+          if (operand.effect === Effect.Freeze) {
+            const operandId =
+              temporaries.get(operand.identifier.id) ?? operand.identifier.id;
+            const lambda = lambdas.get(operandId);
+            if (
+              lambda !== undefined &&
+              lambda.dependencies.some(
+                (place) =>
+                  isMutableEffect(place.effect, place.loc) &&
+                  !isRefValueType(place.identifier) &&
+                  !isUseRefType(place.identifier)
+              )
+            ) {
+              errors.pushErrorDetail(
+                new CompilerErrorDetail({
+                  codeframe: null,
+                  description: null,
+                  loc: typeof operand.loc !== "symbol" ? operand.loc : null,
+                  reason:
+                    "Cannot use a mutable function where an immutable value is expected",
+                  severity: ErrorSeverity.InvalidInput,
+                })
+              );
+            }
+          }
+        }
+      }
+    }
+  }
+  if (errors.hasErrors()) {
+    throw errors;
+  }
+}

compiler/forget/src/HIR/index.ts

@@ -18,6 +18,7 @@ export {
 export { mergeConsecutiveBlocks } from "./MergeConsecutiveBlocks";
 export { printFunction, printHIR } from "./PrintHIR";
 export { validateConsistentIdentifiers } from "./ValidateConsistentIdentifiers";
+export { validateFrozenLambdas } from "./ValidateFrozenLambdas";
 export { validateHooksUsage } from "./ValidateHooksUsage";
 export { validateNoRefAccessInRender } from "./ValidateNoRefAccesInRender";
 export { validateTerminalSuccessors } from "./ValidateTerminalSuccessors";

compiler/forget/src/__tests__/fixtures/compiler/capture-func-passed-to-jsx.expect.md

@@ -1,75 +0,0 @@
-
-## Input
-
-```javascript
-function component(a, b) {
-  let y = { b };
-  let z = { a };
-  let x = function () {
-    z.a = 2;
-    y.b;
-  };
-  let t = <Foo x={x}></Foo>;
-  mutate(x); // x should be frozen here
-  return t;
-}
-
-```
-
-## Code
-
-```javascript
-import { unstable_useMemoCache as useMemoCache } from "react";
-function component(a, b) {
-  const $ = useMemoCache(9);
-  const c_0 = $[0] !== b;
-  let t0;
-  if (c_0) {
-    t0 = { b };
-    $[0] = b;
-    $[1] = t0;
-  } else {
-    t0 = $[1];
-  }
-  const y = t0;
-  const c_2 = $[2] !== a;
-  let t1;
-  if (c_2) {
-    t1 = { a };
-    $[2] = a;
-    $[3] = t1;
-  } else {
-    t1 = $[3];
-  }
-  const z = t1;
-  const c_4 = $[4] !== z.a;
-  const c_5 = $[5] !== y.b;
-  let t2;
-  if (c_4 || c_5) {
-    t2 = function () {
-      z.a = 2;
-      y.b;
-    };
-    $[4] = z.a;
-    $[5] = y.b;
-    $[6] = t2;
-  } else {
-    t2 = $[6];
-  }
-  const x = t2;
-  const c_7 = $[7] !== x;
-  let t3;
-  if (c_7) {
-    t3 = <Foo x={x} />;
-    $[7] = x;
-    $[8] = t3;
-  } else {
-    t3 = $[8];
-  }
-  const t = t3;
-  mutate(x);
-  return t;
-}
-
-```
-      

compiler/forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-and-local-variables-with-default.expect.md

@@ -17,7 +17,7 @@ function Component(props) {
     if (!comments.length) {
       return;
     }
-    log(comments.length);
+    console.log(comments.length);
   };
   allUrls.push(...urls);
   return <Media media={media} onClick={onClick} />;
@@ -38,7 +38,7 @@ function Component(props) {
   if (c_0) {
     const allUrls = [];
 
-    const { media: t0, comments: t2, urls: t81 } = post;
+    const { media: t0, comments: t2, urls: t82 } = post;
     const c_3 = $[3] !== t0;
     let t1;
     if (c_3) {
@@ -59,7 +59,7 @@ function Component(props) {
       t3 = $[6];
     }
     const comments = t3;
-    const urls = t81 === undefined ? [] : t81;
+    const urls = t82 === undefined ? [] : t82;
     const c_7 = $[7] !== comments.length;
     let t4;
     if (c_7) {
@@ -67,7 +67,7 @@ function Component(props) {
         if (!comments.length) {
           return;
         }
-        log(comments.length);
+        console.log(comments.length);
       };
       $[7] = comments.length;
       $[8] = t4;

compiler/forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-and-local-variables-with-default.js

@@ -13,7 +13,7 @@ function Component(props) {
     if (!comments.length) {
       return;
     }
-    log(comments.length);
+    console.log(comments.length);
   };
   allUrls.push(...urls);
   return <Media media={media} onClick={onClick} />;

compiler/forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md

@@ -17,7 +17,7 @@ function Component(props) {
     if (!comments.length) {
       return;
     }
-    log(comments.length);
+    console.log(comments.length);
   };
   allUrls.push(...urls);
   return <Media media={media} onClick={onClick} />;
@@ -38,8 +38,8 @@ function Component(props) {
   if (c_0) {
     const allUrls = [];
 
-    const { media: t83, comments, urls } = post;
-    media = t83;
+    const { media: t85, comments, urls } = post;
+    media = t85;
     const c_3 = $[3] !== comments.length;
     let t0;
     if (c_3) {
@@ -47,7 +47,7 @@ function Component(props) {
         if (!comments.length) {
           return;
         }
-        log(comments.length);
+        console.log(comments.length);
       };
       $[3] = comments.length;
       $[4] = t0;

compiler/forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.js

@@ -13,7 +13,7 @@ function Component(props) {
     if (!comments.length) {
       return;
     }
-    log(comments.length);
+    console.log(comments.length);
   };
   allUrls.push(...urls);
   return <Media media={media} onClick={onClick} />;

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-capture-func-passed-to-jsx.expect.md

@@ -0,0 +1,26 @@
+
+## Input
+
+```javascript
+function component(a, b) {
+  let y = { b };
+  let z = { a };
+  let x = function () {
+    z.a = 2;
+    y.b;
+  };
+  let t = <Foo x={x}></Foo>;
+  mutate(x); // x should be frozen here
+  return t;
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidInput: Cannot use a mutable function where an immutable value is expected (8:8)
+```
+          
+      

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-freeze-mutable-lambda-mutate-local.expect.md

@@ -0,0 +1,24 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const x = {};
+  const onChange = (e) => {
+    // INVALID! should use copy-on-write and pass the new value
+    x.value = e.target.value;
+    setX(x);
+  };
+  return <input value={x.value} onChange={onChange} />;
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidInput: Cannot use a mutable function where an immutable value is expected (8:8)
+```
+          
+      

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-freeze-mutable-lambda-mutate-local.js

@@ -0,0 +1,9 @@
+function Component(props) {
+  const x = {};
+  const onChange = (e) => {
+    // INVALID! should use copy-on-write and pass the new value
+    x.value = e.target.value;
+    setX(x);
+  };
+  return <input value={x.value} onChange={onChange} />;
+}

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-freeze-mutable-lambda-reassign-local.expect.md

@@ -0,0 +1,22 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  let x = "";
+  const onChange = (e) => {
+    x = e.target.value;
+  };
+  return <input value={x} onChange={onChange} />;
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidInput: Cannot use a mutable function where an immutable value is expected (6:6)
+```
+          
+      

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-freeze-mutable-lambda-reassign-local.js

@@ -0,0 +1,7 @@
+function Component(props) {
+  let x = "";
+  const onChange = (e) => {
+    x = e.target.value;
+  };
+  return <input value={x} onChange={onChange} />;
+}