[compiler] HIR-based FlattenScopesWithHooksOrUse

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -71,6 +71,7 @@ import {
 import { alignMethodCallScopes } from "../ReactiveScopes/AlignMethodCallScopes";
 import { alignReactiveScopesToBlockScopesHIR } from "../ReactiveScopes/AlignReactiveScopesToBlockScopesHIR";
 import { flattenReactiveLoopsHIR } from "../ReactiveScopes/FlattenReactiveLoopsHIR";
+import { flattenScopesWithHooksOrUseHIR } from "../ReactiveScopes/FlattenScopesWithHooksOrUseHIR";
 import { pruneAlwaysInvalidatingScopes } from "../ReactiveScopes/PruneAlwaysInvalidatingScopes";
 import pruneInitializationDependencies from "../ReactiveScopes/PruneInitializationDependencies";
 import { stabilizeBlockIds } from "../ReactiveScopes/StabilizeBlockIds";
@@ -289,6 +290,13 @@ function* runWithEnvironment(
       name: "FlattenReactiveLoopsHIR",
       value: hir,
     });
+
+    flattenScopesWithHooksOrUseHIR(hir);
+    yield log({
+      kind: "hir",
+      name: "FlattenScopesWithHooksOrUseHIR",
+      value: hir,
+    });
   }
 
   const reactiveFunction = buildReactiveFunction(hir);
@@ -335,17 +343,17 @@ function* runWithEnvironment(
       name: "FlattenReactiveLoops",
       value: reactiveFunction,
     });
+
+    flattenScopesWithHooksOrUse(reactiveFunction);
+    yield log({
+      kind: "reactive",
+      name: "FlattenScopesWithHooks",
+      value: reactiveFunction,
+    });
   }
 
   assertScopeInstructionsWithinScopes(reactiveFunction);
 
-  flattenScopesWithHooksOrUse(reactiveFunction);
-  yield log({
-    kind: "reactive",
-    name: "FlattenScopesWithHooks",
-    value: reactiveFunction,
-  });
-
   propagateScopeDependencies(reactiveFunction);
   yield log({
     kind: "reactive",

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts

@@ -0,0 +1,112 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { CompilerError } from "..";
+import {
+  BlockId,
+  HIRFunction,
+  LabelTerminal,
+  PrunedScopeTerminal,
+  ReactiveScope,
+  getHookKind,
+  isUseOperator,
+} from "../HIR";
+import { retainWhere } from "../Utils/utils";
+
+/**
+ * For simplicity the majority of compiler passes do not treat hooks specially. However, hooks are different
+ * from regular functions in two key ways:
+ * - They can introduce reactivity even when their arguments are non-reactive (accounted for in InferReactivePlaces)
+ * - They cannot be called conditionally
+ *
+ * The `use` operator is similar:
+ * - It can access context, and therefore introduce reactivity
+ * - It can be called conditionally, but _it must be called if the component needs the return value_. This is because
+ *   React uses the fact that use was called to remember that the component needs the value, and that changes to the
+ *   input should invalidate the component itself.
+ *
+ * This pass accounts for the "can't call conditionally" aspect of both hooks and use. Though the reasoning is slightly
+ * different for reach, the result is that we can't memoize scopes that call hooks or use since this would make them
+ * called conditionally in the output.
+ *
+ * The pass finds and removes any scopes that transitively contain a hook or use 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 flattenScopesWithHooksOrUseHIR(fn: HIRFunction): void {
+  const activeScopes: Array<{ block: BlockId; scope: ReactiveScope }> = [];
+  const prune: Array<BlockId> = [];
+
+  for (const [, block] of fn.body.blocks) {
+    const firstId = block.instructions[0]?.id ?? block.terminal.id;
+    retainWhere(activeScopes, (current) => current.scope.range.end > firstId);
+
+    for (const instr of block.instructions) {
+      const { value } = instr;
+      switch (value.kind) {
+        case "MethodCall":
+        case "CallExpression": {
+          const callee =
+            value.kind === "MethodCall" ? value.property : value.callee;
+          if (
+            getHookKind(fn.env, callee.identifier) != null ||
+            isUseOperator(callee.identifier)
+          ) {
+            prune.push(...activeScopes.map((entry) => entry.block));
+            activeScopes.length = 0;
+          }
+        }
+      }
+    }
+    if (block.terminal.kind === "scope") {
+      activeScopes.push({
+        block: block.id,
+        scope: block.terminal.scope,
+      });
+    }
+  }
+
+  for (const id of prune) {
+    const block = fn.body.blocks.get(id)!;
+    const terminal = block.terminal;
+    CompilerError.invariant(terminal.kind === "scope", {
+      reason: `Expected block to have a scope terminal`,
+      description: `Expected block bb${block.id} to end in a scope terminal`,
+      loc: terminal.loc,
+    });
+    const body = fn.body.blocks.get(terminal.block)!;
+    if (
+      body.instructions.length === 1 &&
+      body.terminal.kind === "goto" &&
+      body.terminal.block === terminal.fallthrough
+    ) {
+      /*
+       * This was a scope just for a hook call, which doesn't need memoization.
+       * flatten it away. We rely on the PrunedUnusedLabel step to do the actual
+       * flattening
+       */
+      block.terminal = {
+        kind: "label",
+        block: terminal.block,
+        fallthrough: terminal.fallthrough,
+        id: terminal.id,
+        loc: terminal.loc,
+      } as LabelTerminal;
+      continue;
+    }
+
+    block.terminal = {
+      kind: "pruned-scope",
+      block: terminal.block,
+      fallthrough: terminal.fallthrough,
+      id: terminal.id,
+      loc: terminal.loc,
+      scope: terminal.scope,
+    } as PrunedScopeTerminal;
+  }
+}