Dont memoize scopes with hook calls

compiler/forget/src/CompilerPipeline.ts

@@ -33,6 +33,7 @@ import {
   pruneUnusedScopes,
   renameVariables,
 } from "./ReactiveScopes";
+import { flattenScopesWithHooks } from "./ReactiveScopes/FlattenScopesWithHooks";
 import { eliminateRedundantPhi, enterSSA, leaveSSA } from "./SSA";
 import { inferTypes } from "./TypeInference";
 import { logHIRFunction, logReactiveFunction } from "./Utils/logger";
@@ -118,6 +119,13 @@ export function* run(
     value: reactiveFunction,
   });
 
+  flattenScopesWithHooks(reactiveFunction);
+  yield log({
+    kind: "reactive",
+    name: "FlattenScopesWithHooks",
+    value: reactiveFunction,
+  });
+
   propagateScopeDependencies(reactiveFunction);
   yield log({
     kind: "reactive",

compiler/forget/src/Inference/InferReferenceEffects.ts

@@ -801,7 +801,7 @@ const HOOKS: Map<string, Hook> = new Map([
 type HookKind = { kind: "State" } | { kind: "Ref" } | { kind: "Custom" };
 type Hook = HookKind & { effectKind: Effect; valueKind: ValueKind };
 
-function parseHookCall(place: Place): Hook | null {
+export function parseHookCall(place: Place): Hook | null {
   const name = place.identifier.name;
   if (name === null || !name.match(/^_?use/)) {
     return null;

compiler/forget/src/ReactiveScopes/FlattenScopesWithHooks.ts

@@ -0,0 +1,66 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {
+  InstructionId,
+  ReactiveFunction,
+  ReactiveScopeBlock,
+  ReactiveStatement,
+  ReactiveValue,
+} from "../HIR";
+import { parseHookCall } from "../Inference/InferReferenceEffects";
+import {
+  ReactiveFunctionTransform,
+  Transformed,
+  visitReactiveFunction,
+} from "./visitors";
+
+/**
+ * Most parts of compilation do not treat hooks specially, because there is no guarantee that custom
+ * hooks obey any particular contract. For example, we can't assume that custom hooks won't modify
+ * their arguments, and we can't assume that hooks return immutable or memoized values. Therefore
+ * earlier passes largely ignore hooks, and may end up creating reactive scopes that contain hook calls.
+ *
+ * This pass then finds and removes any scopes that transitively contain a hook call. By running all
+ * the reactive scope inference first, agnostic of hooks, we know that the reactive scopes accurately
+ * describe the set of values which "construct together", and remove _all_ that memoization in order
+ * to ensure the hook call does not inadvertently become conditional.
+ */
+export function flattenScopesWithHooks(fn: ReactiveFunction): void {
+  visitReactiveFunction(fn, new Transform(), { hasHook: false });
+}
+
+type State = { hasHook: boolean };
+
+class Transform extends ReactiveFunctionTransform<State> {
+  override transformScope(
+    scope: ReactiveScopeBlock,
+    outerState: State
+  ): Transformed<ReactiveStatement> {
+    const innerState: State = { hasHook: false };
+    this.visitScope(scope, innerState);
+    outerState.hasHook ||= innerState.hasHook;
+    if (innerState.hasHook) {
+      return { kind: "replace-many", value: scope.instructions };
+    } else {
+      return { kind: "keep" };
+    }
+  }
+
+  override visitValue(
+    id: InstructionId,
+    value: ReactiveValue,
+    state: State
+  ): void {
+    if (value.kind === "CallExpression") {
+      const hook = parseHookCall(value.callee);
+      if (hook !== null) {
+        state.hasHook = true;
+      }
+    }
+  }
+}

compiler/forget/src/ReactiveScopes/InferReactiveIdentifiers.ts

@@ -12,6 +12,7 @@ import {
   ReactiveInstruction,
   ReactiveScope,
 } from "../HIR/HIR";
+import { parseHookCall } from "../Inference/InferReferenceEffects";
 import {
   eachReactiveValueOperand,
   ReactiveFunctionVisitor,
@@ -40,6 +41,17 @@ class Environment extends ReactiveFunctionVisitor<IdentifierReactivity> {
         break;
       }
     }
+    if (!hasReactiveInput && instr.value.kind === "CallExpression") {
+      // Hooks cannot be memoized. Even if they do not accept any reactive inputs,
+      // they are not guaranteed to memoize their return value, and their result
+      // must be assumed to be reactive.
+      // TODO: use types or an opt-in registry of custom hook information to
+      // allow treating safe hooks as non-reactive.
+      const hook = parseHookCall(instr.value.callee);
+      if (hook !== null) {
+        hasReactiveInput = true;
+      }
+    }
     reactivityMap.set(lval.place.identifier, hasReactiveInput);
 
     if (hasReactiveInput) {

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

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

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

@@ -0,0 +1,44 @@
+
+## Input
+
+```javascript
+function component() {
+  let [x, setX] = useState(0);
+  const handler = (event) => setX(event.target.value);
+  return <input onChange={handler} value={x} />;
+}
+
+```
+
+## Code
+
+```javascript
+function component() {
+  const $ = React.unstable_useMemoCache();
+  const x = useState(0)[0];
+  const setX = useState(0)[1];
+  const c_0 = $[0] !== setX;
+  let handler;
+  if (c_0) {
+    handler = (event) => setX(event.target.value);
+    $[0] = setX;
+    $[1] = handler;
+  } else {
+    handler = $[1];
+  }
+  const c_2 = $[2] !== handler;
+  const c_3 = $[3] !== x;
+  let t0;
+  if (c_2 || c_3) {
+    t0 = <input onChange={handler} value={x}></input>;
+    $[2] = handler;
+    $[3] = x;
+    $[4] = t0;
+  } else {
+    t0 = $[4];
+  }
+  return t0;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/controlled-input.js

@@ -0,0 +1,5 @@
+function component() {
+  let [x, setX] = useState(0);
+  const handler = (event) => setX(event.target.value);
+  return <input onChange={handler} value={x} />;
+}

compiler/forget/src/__tests__/fixtures/hir/hook-call.expect.md

@@ -34,22 +34,18 @@ function Component(props) {
   } else {
     x = $[0];
   }
-  let y;
-  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
-    y = useFreeze(x);
-    $[1] = y;
-  } else {
-    y = $[1];
-  }
+  const y = useFreeze(x);
   foo(y, x);
+  const c_1 = $[1] !== y;
   let t0;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+  if (c_1) {
     t0 = (
       <Component>
         {x}
         {y}
       </Component>
     );
+    $[1] = y;
     $[2] = t0;
   } else {
     t0 = $[2];

compiler/forget/src/__tests__/fixtures/hir/nested-scopes-hook-call.expect.md

@@ -0,0 +1,27 @@
+
+## Input
+
+```javascript
+function component(props) {
+  let x = [];
+  let y = [];
+  y.push(useHook(props.foo));
+  x.push(y);
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+function component(props) {
+  const x = [];
+  const y = [];
+  y.push(useHook(props.foo));
+  x.push(y);
+  return x;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/hir/nested-scopes-hook-call.js

@@ -0,0 +1,7 @@
+function component(props) {
+  let x = [];
+  let y = [];
+  y.push(useHook(props.foo));
+  x.push(y);
+  return x;
+}

compiler/forget/src/__tests__/fixtures/hir/optional-member-expression.expect.md

@@ -27,15 +27,8 @@ function Foo(props) {
     x = $[1];
   }
   const y = x?.b;
-  const c_2 = $[2] !== y;
-  let z;
-  if (c_2) {
-    z = useBar(y);
-    $[2] = y;
-    $[3] = z;
-  } else {
-    z = $[3];
-  }
+
+  const z = useBar(y);
   return z;
 }
 

compiler/forget/src/__tests__/fixtures/hir/template-literal.expect.md

@@ -25,17 +25,7 @@ function componentA(props) {
 }
 
 function componentB(props) {
-  const $ = React.unstable_useMemoCache();
-  const t0 = `hello ${props.a}`;
-  const c_0 = $[0] !== t0;
-  let x;
-  if (c_0) {
-    x = useFoo(t0);
-    $[0] = t0;
-    $[1] = x;
-  } else {
-    x = $[1];
-  }
+  const x = useFoo(`hello ${props.a}`);
   return x;
 }