Treat setState as non-reactive

Updates PruneNonReactiveDependencies to treat setState functions as 
non-reactive, since we know they have a stable identity. This is based on type 
inference and our recently added definitions for useState and its return type, 
so it's conservative and will only work when our inference can prove that the 
scope dependency has the SetState type. 

Note that this approach is simple and has limitations, notably the fact that the 
setState is non-reactive doesn't propagate. But it's simple, trivially correct, 
and already improves codegen somewhat, so i figured it's worth landing for now. 

## Test Plan 

Tested on internal app

compiler/forget/src/HIR/HIR.ts

@@ -987,6 +987,14 @@ export function isUseRefType(id: Identifier): boolean {
   return id.type.kind === "Object" && id.type.shapeId === "BuiltInUseRefId";
 }
 
+export function isUseStateType(id: Identifier): boolean {
+  return id.type.kind === "Object" && id.type.shapeId === "BuiltInUseState";
+}
+
+export function isSetStateType(id: Identifier): boolean {
+  return id.type.kind === "Function" && id.type.shapeId === "BuiltInSetState";
+}
+
 export function getHookKind(env: Environment, id: Identifier): HookKind | null {
   const idType = id.type;
   if (idType.kind === "Function") {

compiler/forget/src/ReactiveScopes/PrintReactiveFunction.ts

@@ -18,6 +18,7 @@ import {
   printIdentifier,
   printInstructionValue,
   printPlace,
+  printType,
 } from "../HIR/PrintHIR";
 import { assertExhaustive } from "../Utils/utils";
 
@@ -59,7 +60,9 @@ export function printReactiveBlock(
 }
 
 function printDependency(dependency: ReactiveScopeDependency): string {
-  const identifier = printIdentifier(dependency.identifier);
+  const identifier =
+    printIdentifier(dependency.identifier) +
+    printType(dependency.identifier.type);
   return `${identifier}${dependency.path.map((prop) => `.${prop}`).join("")}`;
 }
 

compiler/forget/src/ReactiveScopes/PruneNonReactiveDependencies.ts

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

compiler/forget/src/__tests__/fixtures/compiler/concise-arrow-expr.expect.md

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

compiler/forget/src/__tests__/fixtures/compiler/controlled-input.expect.md

@@ -15,28 +15,26 @@ function component() {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function component() {
-  const $ = useMemoCache(5);
+  const $ = useMemoCache(4);
   const [x, setX] = useState(0);
-  const c_0 = $[0] !== setX;
   let t0;
-  if (c_0) {
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = (event) => setX(event.target.value);
-    $[0] = setX;
-    $[1] = t0;
+    $[0] = t0;
   } else {
-    t0 = $[1];
+    t0 = $[0];
   }
   const handler = t0;
-  const c_2 = $[2] !== handler;
-  const c_3 = $[3] !== x;
+  const c_1 = $[1] !== handler;
+  const c_2 = $[2] !== x;
   let t1;
-  if (c_2 || c_3) {
+  if (c_1 || c_2) {
     t1 = <input onChange={handler} value={x} />;
-    $[2] = handler;
-    $[3] = x;
-    $[4] = t1;
+    $[1] = handler;
+    $[2] = x;
+    $[3] = t1;
   } else {
-    t1 = $[4];
+    t1 = $[3];
   }
   return t1;
 }

compiler/forget/src/__tests__/fixtures/compiler/disable-jsx-memoization.expect.md

@@ -22,18 +22,16 @@ function Component(props) {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react"; // @memoizeJsxElements false
 function Component(props) {
-  const $ = useMemoCache(2);
+  const $ = useMemoCache(1);
   const [name, setName] = useState(null);
-  const c_0 = $[0] !== setName;
   let t0;
-  if (c_0) {
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = function (e) {
       setName(e.target.value);
     };
-    $[0] = setName;
-    $[1] = t0;
+    $[0] = t0;
   } else {
-    t0 = $[1];
+    t0 = $[0];
   }
   const onChange = t0;
   return (

compiler/forget/src/__tests__/fixtures/compiler/inadvertent-mutability-readonly-lambda.expect.md

@@ -23,29 +23,27 @@ function Component(props) {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Component(props) {
-  const $ = useMemoCache(4);
+  const $ = useMemoCache(3);
   const [value, setValue] = useState(null);
-  const c_0 = $[0] !== setValue;
   let t0;
-  if (c_0) {
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = (e) => setValue((value_0) => value_0 + e.target.value);
-    $[0] = setValue;
-    $[1] = t0;
+    $[0] = t0;
   } else {
-    t0 = $[1];
+    t0 = $[0];
   }
   const onChange = t0;
 
   useOtherHook();
-  const c_2 = $[2] !== onChange;
+  const c_1 = $[1] !== onChange;
   let x;
-  if (c_2) {
+  if (c_1) {
     x = {};
     foo(x, onChange);
-    $[2] = onChange;
-    $[3] = x;
+    $[1] = onChange;
+    $[2] = x;
   } else {
-    x = $[3];
+    x = $[2];
   }
   return x;
 }

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

@@ -19,7 +19,7 @@ function Component(props) {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Component(props) {
-  const $ = useMemoCache(7);
+  const $ = useMemoCache(6);
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = { value: "" };
@@ -29,31 +29,29 @@ function Component(props) {
   }
   const [x, setX] = useState(t0);
   const c_1 = $[1] !== x;
-  const c_2 = $[2] !== setX;
   let t1;
-  if (c_1 || c_2) {
+  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] = setX;
-    $[3] = t1;
+    $[2] = t1;
   } else {
-    t1 = $[3];
+    t1 = $[2];
   }
   const onChange = t1;
-  const c_4 = $[4] !== x.value;
-  const c_5 = $[5] !== onChange;
+  const c_3 = $[3] !== x.value;
+  const c_4 = $[4] !== onChange;
   let t2;
-  if (c_4 || c_5) {
+  if (c_3 || c_4) {
     t2 = <input value={x.value} onChange={onChange} />;
-    $[4] = x.value;
-    $[5] = onChange;
-    $[6] = t2;
+    $[3] = x.value;
+    $[4] = onChange;
+    $[5] = t2;
   } else {
-    t2 = $[6];
+    t2 = $[5];
   }
   return t2;
 }

compiler/forget/src/__tests__/fixtures/compiler/nested-function-shadowed-identifiers.expect.md

@@ -20,31 +20,29 @@ function Component(props) {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function Component(props) {
-  const $ = useMemoCache(5);
+  const $ = useMemoCache(4);
   const [x, setX] = useState(null);
-  const c_0 = $[0] !== setX;
   let t0;
-  if (c_0) {
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = (e) => {
       let x_0 = null; // intentionally shadow the original x
       setX((currentX) => currentX + x_0); // intentionally refer to shadowed x
     };
-    $[0] = setX;
-    $[1] = t0;
+    $[0] = t0;
   } else {
-    t0 = $[1];
+    t0 = $[0];
   }
   const onChange = t0;
-  const c_2 = $[2] !== x;
-  const c_3 = $[3] !== onChange;
+  const c_1 = $[1] !== x;
+  const c_2 = $[2] !== onChange;
   let t1;
-  if (c_2 || c_3) {
+  if (c_1 || c_2) {
     t1 = <input value={x} onChange={onChange} />;
-    $[2] = x;
-    $[3] = onChange;
-    $[4] = t1;
+    $[1] = x;
+    $[2] = onChange;
+    $[3] = t1;
   } else {
-    t1 = $[4];
+    t1 = $[3];
   }
   return t1;
 }

compiler/forget/src/__tests__/fixtures/compiler/use-callback-simple.expect.md

@@ -16,28 +16,26 @@ function component() {
 ```javascript
 import { unstable_useMemoCache as useMemoCache } from "react";
 function component() {
-  const $ = useMemoCache(5);
+  const $ = useMemoCache(4);
   const [count, setCount] = useState(0);
-  const c_0 = $[0] !== setCount;
-  const c_1 = $[1] !== count;
+  const c_0 = $[0] !== count;
   let t0;
-  if (c_0 || c_1) {
+  if (c_0) {
     t0 = () => setCount(count + 1);
-    $[0] = setCount;
-    $[1] = count;
-    $[2] = t0;
+    $[0] = count;
+    $[1] = t0;
   } else {
-    t0 = $[2];
+    t0 = $[1];
   }
   const increment = t0;
-  const c_3 = $[3] !== increment;
+  const c_2 = $[2] !== increment;
   let t1;
-  if (c_3) {
+  if (c_2) {
     t1 = <Foo onClick={increment} />;
-    $[3] = increment;
-    $[4] = t1;
+    $[2] = increment;
+    $[3] = t1;
   } else {
-    t1 = $[4];
+    t1 = $[3];
   }
   return t1;
 }