@samitouri / QOS-React / commits / b5b9daa44d

[compiler] HIR-based FlattenScopesWithHooksOrUse

Per title, implements an HIR-based version of FlattenScopesWithHooksOrUse as part of our push to use HIR everywhere. This is the last pass to migrate before PropagateScopeDeps, which is blocking the fix for `bug.invalid-hoisting-functionexpr`, ie where we can infer incorrect dependencies for function expressions if the dependencies are accessed conditionally. ghstack-source-id: 05c6e26b3b7a3b1c3e106a37053f88ac3c72caf5 Pull Request resolved: https://github.com/facebook/react/pull/29840

Joe Savona committed Jun 10, 2024 at 14:14 UTC b5b9daa44d3cd5849d2a48da3c1db71410d9e294
2 files changed +127 -7
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+15 -7
@@ -71,6 +71,7 @@ import {
71 import { alignMethodCallScopes } from "../ReactiveScopes/AlignMethodCallScopes";
72 import { alignReactiveScopesToBlockScopesHIR } from "../ReactiveScopes/AlignReactiveScopesToBlockScopesHIR";
73 import { flattenReactiveLoopsHIR } from "../ReactiveScopes/FlattenReactiveLoopsHIR";
74 +import { flattenScopesWithHooksOrUseHIR } from "../ReactiveScopes/FlattenScopesWithHooksOrUseHIR";
75 import { pruneAlwaysInvalidatingScopes } from "../ReactiveScopes/PruneAlwaysInvalidatingScopes";
76 import pruneInitializationDependencies from "../ReactiveScopes/PruneInitializationDependencies";
77 import { stabilizeBlockIds } from "../ReactiveScopes/StabilizeBlockIds";
@@ -289,6 +290,13 @@ function* runWithEnvironment(
290 name: "FlattenReactiveLoopsHIR",
291 value: hir,
292 });
293 +
294 + flattenScopesWithHooksOrUseHIR(hir);
295 + yield log({
296 + kind: "hir",
297 + name: "FlattenScopesWithHooksOrUseHIR",
298 + value: hir,
299 + });
300 }
301
302 const reactiveFunction = buildReactiveFunction(hir);
@@ -335,17 +343,17 @@ function* runWithEnvironment(
343 name: "FlattenReactiveLoops",
344 value: reactiveFunction,
345 });
346 +
347 + flattenScopesWithHooksOrUse(reactiveFunction);
348 + yield log({
349 + kind: "reactive",
350 + name: "FlattenScopesWithHooks",
351 + value: reactiveFunction,
352 + });
353 }
354
355 assertScopeInstructionsWithinScopes(reactiveFunction);
356
342 - flattenScopesWithHooksOrUse(reactiveFunction);
343 - yield log({
344 - kind: "reactive",
345 - name: "FlattenScopesWithHooks",
346 - value: reactiveFunction,
347 - });
348 -
357 propagateScopeDependencies(reactiveFunction);
358 yield log({
359 kind: "reactive",
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts new
+112
@@ -0,0 +1,112 @@
1 +/**
2 + * Copyright (c) Meta Platforms, Inc. and affiliates.
3 + *
4 + * This source code is licensed under the MIT license found in the
5 + * LICENSE file in the root directory of this source tree.
6 + */
7 +
8 +import { CompilerError } from "..";
9 +import {
10 + BlockId,
11 + HIRFunction,
12 + LabelTerminal,
13 + PrunedScopeTerminal,
14 + ReactiveScope,
15 + getHookKind,
16 + isUseOperator,
17 +} from "../HIR";
18 +import { retainWhere } from "../Utils/utils";
19 +
20 +/**
21 + * For simplicity the majority of compiler passes do not treat hooks specially. However, hooks are different
22 + * from regular functions in two key ways:
23 + * - They can introduce reactivity even when their arguments are non-reactive (accounted for in InferReactivePlaces)
24 + * - They cannot be called conditionally
25 + *
26 + * The `use` operator is similar:
27 + * - It can access context, and therefore introduce reactivity
28 + * - It can be called conditionally, but _it must be called if the component needs the return value_. This is because
29 + * React uses the fact that use was called to remember that the component needs the value, and that changes to the
30 + * input should invalidate the component itself.
31 + *
32 + * This pass accounts for the "can't call conditionally" aspect of both hooks and use. Though the reasoning is slightly
33 + * different for reach, the result is that we can't memoize scopes that call hooks or use since this would make them
34 + * called conditionally in the output.
35 + *
36 + * The pass finds and removes any scopes that transitively contain a hook or use call. By running all
37 + * the reactive scope inference first, agnostic of hooks, we know that the reactive scopes accurately
38 + * describe the set of values which "construct together", and remove _all_ that memoization in order
39 + * to ensure the hook call does not inadvertently become conditional.
40 + */
41 +export function flattenScopesWithHooksOrUseHIR(fn: HIRFunction): void {
42 + const activeScopes: Array<{ block: BlockId; scope: ReactiveScope }> = [];
43 + const prune: Array<BlockId> = [];
44 +
45 + for (const [, block] of fn.body.blocks) {
46 + const firstId = block.instructions[0]?.id ?? block.terminal.id;
47 + retainWhere(activeScopes, (current) => current.scope.range.end > firstId);
48 +
49 + for (const instr of block.instructions) {
50 + const { value } = instr;
51 + switch (value.kind) {
52 + case "MethodCall":
53 + case "CallExpression": {
54 + const callee =
55 + value.kind === "MethodCall" ? value.property : value.callee;
56 + if (
57 + getHookKind(fn.env, callee.identifier) != null ||
58 + isUseOperator(callee.identifier)
59 + ) {
60 + prune.push(...activeScopes.map((entry) => entry.block));
61 + activeScopes.length = 0;
62 + }
63 + }
64 + }
65 + }
66 + if (block.terminal.kind === "scope") {
67 + activeScopes.push({
68 + block: block.id,
69 + scope: block.terminal.scope,
70 + });
71 + }
72 + }
73 +
74 + for (const id of prune) {
75 + const block = fn.body.blocks.get(id)!;
76 + const terminal = block.terminal;
77 + CompilerError.invariant(terminal.kind === "scope", {
78 + reason: `Expected block to have a scope terminal`,
79 + description: `Expected block bb${block.id} to end in a scope terminal`,
80 + loc: terminal.loc,
81 + });
82 + const body = fn.body.blocks.get(terminal.block)!;
83 + if (
84 + body.instructions.length === 1 &&
85 + body.terminal.kind === "goto" &&
86 + body.terminal.block === terminal.fallthrough
87 + ) {
88 + /*
89 + * This was a scope just for a hook call, which doesn't need memoization.
90 + * flatten it away. We rely on the PrunedUnusedLabel step to do the actual
91 + * flattening
92 + */
93 + block.terminal = {
94 + kind: "label",
95 + block: terminal.block,
96 + fallthrough: terminal.fallthrough,
97 + id: terminal.id,
98 + loc: terminal.loc,
99 + } as LabelTerminal;
100 + continue;
101 + }
102 +
103 + block.terminal = {
104 + kind: "pruned-scope",
105 + block: terminal.block,
106 + fallthrough: terminal.fallthrough,
107 + id: terminal.id,
108 + loc: terminal.loc,
109 + scope: terminal.scope,
110 + } as PrunedScopeTerminal;
111 + }
112 +}