Integrate mutable context identifier inference, improve frozen function

inference 

Integrates the new inference to detect definitive mutations of context 
variables. This allows us to more precisely infer mutable function expressions 
and validate that they aren't passed where frozen values are expected.

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

@@ -16,7 +16,6 @@ import {
   HIRFunction,
   IdentifierId,
   Place,
-  isMutableEffect,
   isRefValueType,
   isUseRefType,
 } from "./HIR";
@@ -106,7 +105,7 @@ function validateOperand(
       lambda !== undefined &&
       lambda.dependencies.some(
         (place) =>
-          isMutableEffect(place.effect, place.loc) &&
+          place.effect === Effect.Mutate &&
           !isRefValueType(place.identifier) &&
           !isUseRefType(place.identifier)
       )

compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts

@@ -18,9 +18,11 @@ import {
   ReactiveScopeDependency,
 } from "../HIR";
 import { constantPropagation } from "../Optimization";
-import { eliminateRedundantPhi, enterSSA } from "../SSA";
+import { inferReactiveScopeVariables } from "../ReactiveScopes";
+import { eliminateRedundantPhi, enterSSA, leaveSSA } from "../SSA";
 import { inferTypes } from "../TypeInference";
 import { logHIRFunction } from "../Utils/logger";
+import { inferMutableContextVariables } from "./InferMutableContextVariables";
 import { inferMutableRanges } from "./InferMutableRanges";
 import inferReferenceEffects from "./InferReferenceEffects";
 
@@ -112,6 +114,9 @@ function lower(func: HIRFunction): void {
   analyseFunctions(func);
   inferReferenceEffects(func, { isFunctionExpression: true });
   inferMutableRanges(func);
+  leaveSSA(func);
+  inferReactiveScopeVariables(func);
+  inferMutableContextVariables(func);
   logHIRFunction("AnalyseFunction (inner)", func);
 }
 
@@ -120,12 +125,15 @@ function infer(
   state: IdentifierState,
   context: Place[]
 ): void {
-  const mutations = new Set(
-    value.loweredFunc.context
-      .filter((dep) => isMutatedOrReassigned(dep.identifier))
-      .map((m) => m.identifier.name)
-      .filter((m) => m !== null) as string[]
-  );
+  const mutations = new Map<string, Effect>();
+  for (const operand of value.loweredFunc.context) {
+    if (
+      isMutatedOrReassigned(operand.identifier) &&
+      operand.identifier.name !== null
+    ) {
+      mutations.set(operand.identifier.name, operand.effect);
+    }
+  }
 
   for (const dep of value.dependencies) {
     let name: string | null = null;
@@ -144,8 +152,11 @@ function infer(
       // could be called, and allows us to help ensure it isn't called during
       // render
       dep.effect = Effect.Capture;
-    } else if (name !== null && mutations.has(name)) {
-      dep.effect = Effect.Capture;
+    } else if (name !== null) {
+      const effect = mutations.get(name);
+      if (effect !== undefined) {
+        dep.effect = effect === Effect.Unknown ? Effect.Capture : effect;
+      }
     }
   }
 
@@ -161,8 +172,9 @@ function infer(
       "context refs should always have a name"
     );
 
-    if (mutations.has(place.identifier.name)) {
-      place.effect = Effect.Capture;
+    const effect = mutations.get(place.identifier.name);
+    if (effect !== undefined) {
+      place.effect = effect === Effect.Unknown ? Effect.Capture : effect;
       value.dependencies.push(place);
     }
   }

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/capturing-function-conditional-capture-mutate.expect.md

@@ -0,0 +1,62 @@
+
+## Input
+
+```javascript
+// @debug
+function component(a, b) {
+  let z = { a };
+  let y = b;
+  let x = function () {
+    if (y) {
+      // we don't know for sure this mutates, so we should assume
+      // that there is no mutation so long as `x` isn't called
+      // during render
+      maybeMutate(z);
+    }
+  };
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react"; // @debug
+function component(a, b) {
+  const $ = useMemoCache(5);
+  const c_0 = $[0] !== a;
+  let t0;
+  if (c_0) {
+    t0 = { a };
+    $[0] = a;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const z = t0;
+  const y = b;
+  const c_2 = $[2] !== y;
+  const c_3 = $[3] !== z;
+  let t1;
+  if (c_2 || c_3) {
+    t1 = function () {
+      if (y) {
+        // we don't know for sure this mutates, so we should assume
+        // that there is no mutation so long as `x` isn't called
+        // during render
+        maybeMutate(z);
+      }
+    };
+    $[2] = y;
+    $[3] = z;
+    $[4] = t1;
+  } else {
+    t1 = $[4];
+  }
+  const x = t1;
+  return x;
+}
+
+```
+      

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/capturing-function-conditional-capture-mutate.js

@@ -1,3 +1,4 @@
+// @debug
 function component(a, b) {
   let z = { a };
   let y = b;

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.invalid-function-expression-mutates-immutable-value.expect.md

@@ -0,0 +1,24 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const [x, setX] = useState({ value: "" });
+  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: InferReferenceEffects: inferred mutation of known immutable value. Found mutation of $40 (frozen) (5:5)
+```
+          
+      

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-capturing-function-conditional-capture-mutate.expect.md

@@ -1,28 +0,0 @@
-
-## Input
-
-```javascript
-function component(a, b) {
-  let z = { a };
-  let y = b;
-  let x = function () {
-    if (y) {
-      // we don't know for sure this mutates, so we should assume
-      // that there is no mutation so long as `x` isn't called
-      // during render
-      maybeMutate(z);
-    }
-  };
-  return x;
-}
-
-```
-
-
-## Error
-
-```
-[ReactForget] InvalidInput: Cannot use a mutable function where an immutable value is expected (12:12)
-```
-          
-      

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-function-expression-captures-value-later-frozen-jsx.expect.md

@@ -1,27 +0,0 @@
-
-## Input
-
-```javascript
-function Component(props) {
-  let x = {};
-  // onChange should be inferred as immutable, because the value
-  // it captures (`x`) is frozen by the time the function is referenced
-  const onChange = (e) => {
-    maybeMutate(x, e.target.value);
-  };
-  if (props.cond) {
-    <div>{x}</div>;
-  }
-  return <Foo value={x} onChange={onChange} />;
-}
-
-```
-
-
-## Error
-
-```
-[ReactForget] InvalidInput: Cannot use a mutable function where an immutable value is expected (11:11)
-```
-          
-      

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expression-captures-value-later-frozen-jsx.expect.md

@@ -0,0 +1,57 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  let x = {};
+  // onChange should be inferred as immutable, because the value
+  // it captures (`x`) is frozen by the time the function is referenced
+  const onChange = (e) => {
+    maybeMutate(x, e.target.value);
+  };
+  if (props.cond) {
+    <div>{x}</div>;
+  }
+  return <Foo value={x} onChange={onChange} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = {};
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const x = t0;
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = (e) => {
+      maybeMutate(x, e.target.value);
+    };
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const onChange = t1;
+  if (props.cond) {
+  }
+  let t2;
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = <Foo value={x} onChange={onChange} />;
+    $[2] = t2;
+  } else {
+    t2 = $[2];
+  }
+  return t2;
+}
+
+```
+      

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-expression-maybe-mutates-hook-return-value.expect.md

@@ -17,11 +17,36 @@ function Component(props) {
 
 ```
 
+## Code
 
-## Error
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(4);
+  const id = useSelectedEntitytId();
+  const c_0 = $[0] !== id;
+  let t0;
+  if (c_0) {
+    t0 = () => {
+      log(id);
+    };
+    $[0] = id;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const onLoad = t0;
+  const c_2 = $[2] !== onLoad;
+  let t1;
+  if (c_2) {
+    t1 = <Foo onLoad={onLoad} />;
+    $[2] = onLoad;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  return t1;
+}
 
 ```
-[ReactForget] InvalidInput: Cannot use a mutable function where an immutable value is expected (11:11)
-```
-          
       

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/todo-invalid-function-expression-mutates-immutable-value.expect.md

@@ -1,60 +0,0 @@
-
-## Input
-
-```javascript
-function Component(props) {
-  const [x, setX] = useState({ value: "" });
-  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} />;
-}
-
-```
-
-## Code
-
-```javascript
-import { unstable_useMemoCache as useMemoCache } from "react";
-function Component(props) {
-  const $ = useMemoCache(6);
-  let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = { value: "" };
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  const [x, setX] = useState(t0);
-  const c_1 = $[1] !== x;
-  let t1;
-  if (c_1) {
-    t1 = (e) => {
-      // INVALID! should use copy-on-write and pass the new value
-      x.value = e.target.value;
-      setX(x);
-    };
-    $[1] = x;
-    $[2] = t1;
-  } else {
-    t1 = $[2];
-  }
-  const onChange = t1;
-  const c_3 = $[3] !== x.value;
-  const c_4 = $[4] !== onChange;
-  let t2;
-  if (c_3 || c_4) {
-    t2 = <input value={x.value} onChange={onChange} />;
-    $[3] = x.value;
-    $[4] = onChange;
-    $[5] = t2;
-  } else {
-    t2 = $[5];
-  }
-  return t2;
-}
-
-```
-