InferReactivePlaces understands setState type

I realized this while working on Forest. When computing the dependencies of a 
reactive scope we can omit setState functions in the general case (exception 
described below). Currently that's implemented in PruneNonReactiveDependencies. 
However, this causes us to miss some optimizations — a value isn't reactive if 
its only dependency is a setState, and that may allow further downstreams values 
to become non-reactive. We lose out on that by only filtering out setStates in 
PruneNonReactiveDependencies — this logic really belongs in InferReactivePlaces. 

So this PR moves the check for setState types to that pass. The updated fixtures 
show that this already uncovers some wins. The _new_ fixtures covers the 
exception. It's possible for a value to be typed as being a setState function, 
but to still be reactive: if its a local that is conditionally assigned 
different setState function values. Currently this test happens to work because 
our phi type inference is incomplete (see #2296). I'm adding the test now though 
to prevent regressions when we fix phi type inference.

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

@@ -15,6 +15,7 @@ import {
   Place,
   computePostDominatorTree,
   getHookKind,
+  isSetStateType,
 } from "../HIR";
 import { PostDominator } from "../HIR/Dominator";
 import {
@@ -191,6 +192,9 @@ export function inferReactivePlaces(fn: HIRFunction): void {
 
         if (hasReactiveInput) {
           for (const lvalue of eachInstructionLValue(instruction)) {
+            if (isSetStateType(lvalue.identifier)) {
+              continue;
+            }
             reactiveIdentifiers.markReactive(lvalue);
           }
 

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonReactiveDependencies.ts

@@ -5,12 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {
-  IdentifierId,
-  ReactiveFunction,
-  ReactiveScopeBlock,
-  isSetStateType,
-} from "../HIR";
+import { IdentifierId, ReactiveFunction, ReactiveScopeBlock } from "../HIR";
 import { collectReactiveIdentifiers } from "./CollectReactiveIdentifiers";
 import { ReactiveFunctionVisitor, visitReactiveFunction } from "./visitors";
 
@@ -34,8 +29,7 @@ class Visitor extends ReactiveFunctionVisitor<ReactiveIdentifiers> {
   ): void {
     this.traverseScope(scope, state);
     for (const dep of scope.scope.dependencies) {
-      const isReactive =
-        state.has(dep.identifier.id) && !isSetStateType(dep.identifier);
+      const isReactive = state.has(dep.identifier.id);
       if (!isReactive) {
         scope.scope.dependencies.delete(dep);
       }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/concise-arrow-expr.expect.md

@@ -15,7 +15,7 @@ function component() {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function component() {
-  const $ = useMemoCache(3);
+  const $ = useMemoCache(2);
   const [x, setX] = useState(0);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -26,12 +26,11 @@ function component() {
   }
   const handler = t0;
   let t1;
-  if ($[1] !== handler) {
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
     t1 = <Foo handler={handler} />;
-    $[1] = handler;
-    $[2] = t1;
+    $[1] = t1;
   } else {
-    t1 = $[2];
+    t1 = $[1];
   }
   return t1;
 }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/controlled-input.expect.md

@@ -22,7 +22,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { useState, unstable_useMemoCache as useMemoCache } from "react";
 function component() {
-  const $ = useMemoCache(4);
+  const $ = useMemoCache(3);
   const [x, setX] = useState(0);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -33,13 +33,12 @@ function component() {
   }
   const handler = t0;
   let t1;
-  if ($[1] !== handler || $[2] !== x) {
+  if ($[1] !== x) {
     t1 = <input onChange={handler} value={x} />;
-    $[1] = handler;
-    $[2] = x;
-    $[3] = t1;
+    $[1] = x;
+    $[2] = t1;
   } else {
-    t1 = $[3];
+    t1 = $[2];
   }
   return t1;
 }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/inadvertent-mutability-readonly-lambda.expect.md

@@ -23,7 +23,7 @@ function Component(props) {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Component(props) {
-  const $ = useMemoCache(3);
+  const $ = useMemoCache(2);
   const [value, setValue] = useState(null);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -36,13 +36,12 @@ function Component(props) {
 
   useOtherHook();
   let x;
-  if ($[1] !== onChange) {
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
     x = {};
     foo(x, onChange);
-    $[1] = onChange;
-    $[2] = x;
+    $[1] = x;
   } else {
-    x = $[2];
+    x = $[1];
   }
   return x;
 }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/nested-function-shadowed-identifiers.expect.md

@@ -26,7 +26,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Component(props) {
-  const $ = useMemoCache(4);
+  const $ = useMemoCache(3);
   const [x, setX] = useState(null);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -39,13 +39,12 @@ function Component(props) {
   }
   const onChange = t0;
   let t1;
-  if ($[1] !== x || $[2] !== onChange) {
+  if ($[1] !== x) {
     t1 = <input value={x} onChange={onChange} />;
     $[1] = x;
-    $[2] = onChange;
-    $[3] = t1;
+    $[2] = t1;
   } else {
-    t1 = $[3];
+    t1 = $[2];
   }
   return t1;
 }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/reactive-control-dependency-phi-setState-type.expect.md

@@ -0,0 +1,111 @@
+
+## Input
+
+```javascript
+import invariant from "invariant";
+import { useState } from "react";
+
+function Component(props) {
+  const [x, setX] = useState(false);
+  const [y, setY] = useState(false);
+  let setState;
+  if (props.cond) {
+    setState = setX;
+  } else {
+    setState = setY;
+  }
+  const setState2 = setState;
+  const stateObject = { setState: setState2 };
+  return (
+    <Foo
+      cond={props.cond}
+      setX={setX}
+      setY={setY}
+      setState={stateObject.setState}
+    />
+  );
+}
+
+function Foo({ cond, setX, setY, setState }) {
+  if (cond) {
+    invariant(setState === setX, "Expected the correct setState function");
+  } else {
+    invariant(setState === setY, "Expected the correct setState function");
+  }
+  return "ok";
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  // TODO: run this function with {cond:true}, {cond: false}
+  params: [{ cond: true }],
+};
+
+```
+
+## Code
+
+```javascript
+import invariant from "invariant";
+import { useState, unstable_useMemoCache as useMemoCache } from "react";
+
+function Component(props) {
+  const $ = useMemoCache(5);
+  const [x, setX] = useState(false);
+  const [y, setY] = useState(false);
+  let setState;
+  if (props.cond) {
+    setState = setX;
+  } else {
+    setState = setY;
+  }
+
+  const setState2 = setState;
+  let t0;
+  if ($[0] !== setState2) {
+    t0 = { setState: setState2 };
+    $[0] = setState2;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const stateObject = t0;
+  let t1;
+  if ($[2] !== props.cond || $[3] !== stateObject.setState) {
+    t1 = (
+      <Foo
+        cond={props.cond}
+        setX={setX}
+        setY={setY}
+        setState={stateObject.setState}
+      />
+    );
+    $[2] = props.cond;
+    $[3] = stateObject.setState;
+    $[4] = t1;
+  } else {
+    t1 = $[4];
+  }
+  return t1;
+}
+
+function Foo(t21) {
+  const { cond, setX, setY, setState } = t21;
+  if (cond) {
+    invariant(setState === setX, "Expected the correct setState function");
+  } else {
+    invariant(setState === setY, "Expected the correct setState function");
+  }
+  return "ok";
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  // TODO: run this function with {cond:true}, {cond: false}
+  params: [{ cond: true }],
+};
+
+```
+      
+### Eval output
+(kind: ok) ok

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/reactive-control-dependency-phi-setState-type.js

@@ -0,0 +1,38 @@
+import invariant from "invariant";
+import { useState } from "react";
+
+function Component(props) {
+  const [x, setX] = useState(false);
+  const [y, setY] = useState(false);
+  let setState;
+  if (props.cond) {
+    setState = setX;
+  } else {
+    setState = setY;
+  }
+  const setState2 = setState;
+  const stateObject = { setState: setState2 };
+  return (
+    <Foo
+      cond={props.cond}
+      setX={setX}
+      setY={setY}
+      setState={stateObject.setState}
+    />
+  );
+}
+
+function Foo({ cond, setX, setY, setState }) {
+  if (cond) {
+    invariant(setState === setX, "Expected the correct setState function");
+  } else {
+    invariant(setState === setY, "Expected the correct setState function");
+  }
+  return "ok";
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  // TODO: run this function with {cond:true}, {cond: false}
+  params: [{ cond: true }],
+};