Memoize hook args (treat as escaping)

Updates `PruneNonEscapingScopes` to consider hook arguments as potentially 
escaping. This is because hook inputs are "owned" by React — for example, 
closures passed to `useEffect`, or a value that is passed to a custom hook and 
which then becomes a memoized input.

compiler/forget/src/HIR/Globals.ts

@@ -169,6 +169,24 @@ const BUILTIN_HOOKS: Array<[string, Hook]> = [
       valueKind: ValueKind.Frozen,
     },
   ],
+  [
+    "useEffect",
+    {
+      kind: "Memo",
+      name: "useEffect",
+      effectKind: Effect.Freeze,
+      valueKind: ValueKind.Frozen,
+    },
+  ],
+  [
+    "useLayoutEffect",
+    {
+      kind: "Memo",
+      name: "useLayoutEffect",
+      effectKind: Effect.Freeze,
+      valueKind: ValueKind.Frozen,
+    },
+  ],
 ];
 
 export type Global = BuiltInType | HookType | PolyType;

compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts

@@ -12,6 +12,7 @@ import {
   Effect,
   IdentifierId,
   InstructionId,
+  isHookType,
   Pattern,
   Place,
   ReactiveFunction,
@@ -23,6 +24,7 @@ import {
   ReactiveValue,
   ScopeId,
 } from "../HIR";
+import { eachInstructionValueOperand } from "../HIR/visitors";
 import { log } from "../Utils/logger";
 import { assertExhaustive } from "../Utils/utils";
 import { getPlaceScope } from "./BuildReactiveBlocks";
@@ -37,8 +39,13 @@ import {
 
 /**
  * This pass prunes reactive scopes that are not necessary to bound downstream computation.
- * Specifically, the pass identifies the set of identifiers which are directly returned by
- * the function and/or transitively aliased by a return value - ie, values that "escape".
+ * Specifically, the pass identifies the set of identifiers which may "escape". Values can
+ * escape in one of two ways:
+ * * They are directly returned by the function and/or transitively aliased by a return
+ *   value.
+ * * They are passed as input to a hook. This is because any value passed to a hook may
+ *   have its referenced ultimately stored by React (ie, be aliased by an external value).
+ *   For example, the closure passed to useEffect escapes.
  *
  * Example to build intuition:
  *
@@ -96,7 +103,8 @@ import {
  *      b. Conditional and logical expressions (and a few others) are conditinally aliased,
  *         depending on whether their result value is aliased.
  *      c. JSX is always unaliased (though its props children may be)
- * 2. The same pass which builds the graph also stores the set of returned identifiers.
+ * 2. The same pass which builds the graph also stores the set of returned identifiers and set of
+ *    identifiers passed as arguments to hooks.
  * 3. We traverse the graph starting from the returned identifiers and mark reachable dependencies
  *    as escaping, based on the combination of the parent node's type and its children (eg a
  *    conditional node with an aliased dep promotes to aliased).
@@ -198,7 +206,7 @@ class State {
 
   identifiers: Map<IdentifierId, IdentifierNode> = new Map();
   scopes: Map<ScopeId, ScopeNode> = new Map();
-  returned: Set<IdentifierId> = new Set();
+  escapingValues: Set<IdentifierId> = new Set();
 
   /**
    * Declare a new identifier, used for function id and params
@@ -302,8 +310,8 @@ function computeMemoizedIdentifiers(state: State): Set<IdentifierId> {
   }
 
   // Walk from the "roots" aka returned identifiers.
-  for (const returned of state.returned) {
-    visit(returned);
+  for (const value of state.escapingValues) {
+    visit(value);
   }
 
   return memoized;
@@ -649,6 +657,13 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
         instruction.lvalue.identifier.id,
         instruction.value.place.identifier.id
       );
+    } else if (instruction.value.kind === "CallExpression") {
+      const callee = instruction.value.callee;
+      if (isHookType(callee.identifier)) {
+        for (const operand of eachInstructionValueOperand(instruction.value)) {
+          state.escapingValues.add(operand.identifier.id);
+        }
+      }
     }
   }
 
@@ -659,7 +674,7 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
     this.traverseTerminal(stmt, state);
 
     if (stmt.terminal.kind === "return" && stmt.terminal.value !== null) {
-      state.returned.add(stmt.terminal.value.identifier.id);
+      state.escapingValues.add(stmt.terminal.value.identifier.id);
     }
   }
 }

compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.expect.md

@@ -0,0 +1,75 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const dispatch = useDispatch();
+  useFreeze(dispatch);
+
+  // onUpdate should be memoized even though it doesn't
+  // flow into the return value
+  const onUpdate = () => {
+    dispatch({ kind: "update" });
+  };
+
+  useEffect(() => {
+    onUpdate();
+  }, [onUpdate]);
+
+  return <div />;
+}
+
+```
+
+## Code
+
+```javascript
+function Component(props) {
+  const $ = React.unstable_useMemoCache(7);
+  const dispatch = useDispatch();
+  useFreeze(dispatch);
+  const c_0 = $[0] !== dispatch;
+  let t0;
+  if (c_0) {
+    t0 = () => {
+      dispatch({ kind: "update" });
+    };
+    $[0] = dispatch;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const onUpdate = t0;
+  const c_2 = $[2] !== onUpdate;
+  let t1;
+  if (c_2) {
+    t1 = () => {
+      onUpdate();
+    };
+    $[2] = onUpdate;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  const c_4 = $[4] !== onUpdate;
+  let t2;
+  if (c_4) {
+    t2 = [onUpdate];
+    $[4] = onUpdate;
+    $[5] = t2;
+  } else {
+    t2 = $[5];
+  }
+  useEffect(t1, t2);
+  let t3;
+  if ($[6] === Symbol.for("react.memo_cache_sentinel")) {
+    t3 = <div />;
+    $[6] = t3;
+  } else {
+    t3 = $[6];
+  }
+  return t3;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.js

@@ -0,0 +1,16 @@
+function Component(props) {
+  const dispatch = useDispatch();
+  useFreeze(dispatch);
+
+  // onUpdate should be memoized even though it doesn't
+  // flow into the return value
+  const onUpdate = () => {
+    dispatch({ kind: "update" });
+  };
+
+  useEffect(() => {
+    onUpdate();
+  }, [onUpdate]);
+
+  return <div />;
+}