From b5b9daa44d3cd5849d2a48da3c1db71410d9e294 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Mon, 10 Jun 2024 14:14:52 -0700 Subject: [PATCH] [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 --- .../src/Entrypoint/Pipeline.ts | 22 ++-- .../FlattenScopesWithHooksOrUseHIR.ts | 112 ++++++++++++++++++ 2 files changed, 127 insertions(+), 7 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts index 97ec36626b67b..6ba3b5dcc3e60 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/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", diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUseHIR.ts new file mode 100644 index 0000000000000..d0c0ada10c467 --- /dev/null +++ b/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 = []; + + 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; + } +}