From 9e4b4a5db9a1db387dfa79c277c9d71a9c6d71a1 Mon Sep 17 00:00:00 2001 From: Mike Vitousek Date: Wed, 29 May 2024 12:29:50 -0700 Subject: [PATCH 1/2] [compiler] Prune dependencies that are only used by useRef or useState Summary: @jmbrown215 recently had an observation that the arguments to useState/useRef are only used when a component renders for the first time, and never afterwards. We can skip more computation that we previously could, with reactive blocks that previously recomputed values when inputs changed now only ever computing them on the first render. [ghstack-poisoned] --- .../src/Entrypoint/Pipeline.ts | 8 + .../PruneInitializationDependencies.ts | 290 ++++++++++++++++++ ...d-other-hook-unpruned-dependency.expect.md | 77 +++++ ...tate-and-other-hook-unpruned-dependency.js | 22 ++ .../useState-pruned-dependency.expect.md | 66 ++++ .../compiler/useState-pruned-dependency.js | 17 + .../useState-unpruned-dependency.expect.md | 78 +++++ .../compiler/useState-unpruned-dependency.js | 22 ++ 8 files changed, 580 insertions(+) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.js create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency.js create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.js 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 1fa755499ed..5f50b7a0f1a 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -91,6 +91,7 @@ import { validatePreservedManualMemoization, validateUseMemo, } from "../Validation"; +import pruneInitializationDependencies from "../ReactiveScopes/PruneInitializationDependencies"; export type CompilerPipelineValue = | { kind: "ast"; name: string; value: CodegenFunction } @@ -373,6 +374,13 @@ function* runWithEnvironment( value: reactiveFunction, }); + pruneInitializationDependencies(reactiveFunction); + yield log({ + kind: "reactive", + name: "PruneInitializationDependencies", + value: reactiveFunction, + }); + propagateEarlyReturns(reactiveFunction); yield log({ kind: "reactive", diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts new file mode 100644 index 00000000000..3c181958bee --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneInitializationDependencies.ts @@ -0,0 +1,290 @@ +/** + * 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 { + Environment, + Identifier, + IdentifierId, + InstructionId, + Place, + ReactiveBlock, + ReactiveFunction, + ReactiveInstruction, + ReactiveScopeBlock, + getHookKind, + isUseRefType, + isUseStateType, +} from "../HIR"; +import { eachCallArgument, eachInstructionLValue } from "../HIR/visitors"; +import DisjointSet from "../Utils/DisjointSet"; +import { assertExhaustive } from "../Utils/utils"; +import { ReactiveFunctionVisitor, visitReactiveFunction } from "./visitors"; + +/** + * This pass is built based on the observation by @jbrown215 that arguments + * to useState and useRef are only used the first time a component is rendered. + * Any subsequent times, the arguments will be evaluated but ignored. In this pass, + * we use this fact to improve the output of the compiler by not recomputing values that + * are only used as arguments (or inputs to arguments to) useState and useRef. + * + * Algorithm: + * We take two passes over the reactive function AST. In the first pass, we gather + * aliases and build relationships between property accesses--the key thing we need + * to do here is to find that, e.g., $0.x and $1 refer to the same value if + * $1 = PropertyLoad $0.x. + * + * In the second pass, we traverse the AST in reverse order and track how each place + * is used. If a place is read from in any Terminal, we mark the place as "Update", meaning + * it is used whenever the component is updated/re-rendered. If a place is read from in + * a useState or useRef hook call, we mark it as "Create", since it is only used when the + * component is created. In other instructions, we propagate the inferred place for the + * instructions lvalues onto any other instructions that are read. + * + * Whenever we finish this reverse pass over a reactive block, we can look at the blocks + * dependencies and see whether the dependencies are used in an "Update" context or only + * in a "Create" context. If a dependency is create-only, then we can remove that dependency + * from the block. + */ + +type CreateUpdate = "Create" | "Update" | "Unknown"; + +type KindMap = Map; + +class Visitor extends ReactiveFunctionVisitor { + map: KindMap = new Map(); + aliases: DisjointSet; + paths: Map>; + env: Environment; + + constructor( + env: Environment, + aliases: DisjointSet, + paths: Map> + ) { + super(); + this.aliases = aliases; + this.paths = paths; + this.env = env; + } + + join(values: Array): CreateUpdate { + function join2(l: CreateUpdate, r: CreateUpdate): CreateUpdate { + if (l === r) { + return l; + } + if (l === "Unknown") { + return r; + } + if (r === "Unknown") { + return l; + } + if (l === "Create") { + return r; + } + if (r === "Create") { + return l; + } + if (r === "Update" || l === "Update") { + return "Update"; + } + assertExhaustive(r, `Unhandled variable kind ${r}`); + } + return values.reduce(join2, "Unknown"); + } + + isCreateOnlyHook(id: Identifier): boolean { + return isUseStateType(id) || isUseRefType(id); + } + + override visitPlace( + _: InstructionId, + place: Place, + state: CreateUpdate + ): void { + this.map.set( + place.identifier.id, + this.join([state, this.map.get(place.identifier.id) ?? "Unknown"]) + ); + } + + override visitBlock(block: ReactiveBlock, state: CreateUpdate): void { + super.visitBlock([...block].reverse(), state); + } + + override visitInstruction(instruction: ReactiveInstruction): void { + let state = this.join( + [...eachInstructionLValue(instruction)].map( + (operand) => this.map.get(operand.identifier.id) ?? "Unknown" + ) + ); + + const visitCallOrMethodNonArgs = (): void => { + switch (instruction.value.kind) { + case "CallExpression": { + this.visitPlace(instruction.id, instruction.value.callee, state); + break; + } + case "MethodCall": { + this.visitPlace(instruction.id, instruction.value.property, state); + this.visitPlace(instruction.id, instruction.value.receiver, state); + break; + } + } + }; + + const isHook = (): boolean => { + let callee = null; + switch (instruction.value.kind) { + case "CallExpression": { + callee = instruction.value.callee.identifier; + break; + } + case "MethodCall": { + callee = instruction.value.property.identifier; + break; + } + } + return callee != null && getHookKind(this.env, callee) != null; + }; + + switch (instruction.value.kind) { + case "CallExpression": + case "MethodCall": { + if ( + instruction.lvalue && + this.isCreateOnlyHook(instruction.lvalue.identifier) + ) { + [...eachCallArgument(instruction.value.args)].forEach((operand) => + this.visitPlace(instruction.id, operand, "Create") + ); + visitCallOrMethodNonArgs(); + } else { + if (isHook()) { + /* + * Values flowing into hooks that aren't create-only should be treated + * as Update. + */ + state = "Update"; + } + this.traverseInstruction(instruction, state); + } + break; + } + default: { + this.traverseInstruction(instruction, state); + } + } + } + + override visitScope(scope: ReactiveScopeBlock): void { + const state = this.join( + [ + ...scope.scope.declarations.keys(), + ...[...scope.scope.reassignments.values()].map((ident) => ident.id), + ].map((id) => this.map.get(id) ?? "Unknown") + ); + super.visitScope(scope, state); + [...scope.scope.dependencies].forEach((ident) => { + let target: undefined | IdentifierId = + this.aliases.find(ident.identifier.id) ?? ident.identifier.id; + ident.path.forEach((key) => { + target &&= this.paths.get(target)?.get(key); + }); + if (target && this.map.get(target) === "Create") { + scope.scope.dependencies.delete(ident); + } + }); + } + + override visitReactiveFunctionValue( + _id: InstructionId, + _dependencies: Array, + fn: ReactiveFunction, + state: CreateUpdate + ): void { + visitReactiveFunction(fn, this, state); + } +} + +export default function pruneInitializationDependencies( + fn: ReactiveFunction +): void { + const [aliases, paths] = getAliases(fn); + visitReactiveFunction(fn, new Visitor(fn.env, aliases, paths), "Update"); +} + +function update( + map: Map>, + key: IdentifierId, + path: string, + value: IdentifierId +): void { + const inner = map.get(key) ?? new Map(); + inner.set(path, value); + map.set(key, inner); +} + +class AliasVisitor extends ReactiveFunctionVisitor { + scopeIdentifiers: DisjointSet = new DisjointSet(); + scopePaths: Map> = new Map(); + + override visitInstruction(instr: ReactiveInstruction): void { + if ( + instr.value.kind === "StoreLocal" || + instr.value.kind === "StoreContext" + ) { + this.scopeIdentifiers.union([ + instr.value.lvalue.place.identifier.id, + instr.value.value.identifier.id, + ]); + } else if ( + instr.value.kind === "LoadLocal" || + instr.value.kind === "LoadContext" + ) { + instr.lvalue && + this.scopeIdentifiers.union([ + instr.lvalue.identifier.id, + instr.value.place.identifier.id, + ]); + } else if (instr.value.kind === "PropertyLoad") { + instr.lvalue && + update( + this.scopePaths, + instr.value.object.identifier.id, + instr.value.property, + instr.lvalue.identifier.id + ); + } else if (instr.value.kind === "PropertyStore") { + update( + this.scopePaths, + instr.value.object.identifier.id, + instr.value.property, + instr.value.value.identifier.id + ); + } + } +} + +function getAliases( + fn: ReactiveFunction +): [DisjointSet, Map>] { + const visitor = new AliasVisitor(); + visitReactiveFunction(fn, visitor, null); + let disjoint = visitor.scopeIdentifiers; + let scopePaths = new Map>(); + for (const [key, value] of visitor.scopePaths) { + for (const [path, id] of value) { + update( + scopePaths, + disjoint.find(key) ?? key, + path, + disjoint.find(id) ?? id + ); + } + } + return [disjoint, scopePaths]; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.expect.md new file mode 100644 index 00000000000..1ed67728bde --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.expect.md @@ -0,0 +1,77 @@ + +## Input + +```javascript +import { useState } from "react"; + +function useOther(x) { + return x; +} + +function Component(props) { + const w = f(props.x); + const z = useOther(w); + const [x, _] = useState(z); + return
{x}
; +} + +function f(x) { + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; +import { useState } from "react"; + +function useOther(x) { + return x; +} + +function Component(props) { + const $ = _c(4); + let t0; + if ($[0] !== props.x) { + t0 = f(props.x); + $[0] = props.x; + $[1] = t0; + } else { + t0 = $[1]; + } + const w = t0; + const z = useOther(w); + const [x] = useState(z); + let t1; + if ($[2] !== x) { + t1 =
{x}
; + $[2] = x; + $[3] = t1; + } else { + t1 = $[3]; + } + return t1; +} + +function f(x) { + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; + +``` + +### Eval output +(kind: ok)
42
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.js new file mode 100644 index 00000000000..5c941c3818f --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-and-other-hook-unpruned-dependency.js @@ -0,0 +1,22 @@ +import { useState } from "react"; + +function useOther(x) { + return x; +} + +function Component(props) { + const w = f(props.x); + const z = useOther(w); + const [x, _] = useState(z); + return
{x}
; +} + +function f(x) { + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency.expect.md new file mode 100644 index 00000000000..c5901e0ae79 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency.expect.md @@ -0,0 +1,66 @@ + +## Input + +```javascript +import { useState } from "react"; + +function Component(props) { + const w = f(props.x); + const [x, _] = useState(w); + return
{x}
; +} + +function f(x) { + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; +import { useState } from "react"; + +function Component(props) { + const $ = _c(3); + let t0; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t0 = f(props.x); + $[0] = t0; + } else { + t0 = $[0]; + } + const w = t0; + const [x] = useState(w); + let t1; + if ($[1] !== x) { + t1 =
{x}
; + $[1] = x; + $[2] = t1; + } else { + t1 = $[2]; + } + return t1; +} + +function f(x) { + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; + +``` + +### Eval output +(kind: ok)
42
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency.js new file mode 100644 index 00000000000..8c57ebf72a4 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency.js @@ -0,0 +1,17 @@ +import { useState } from "react"; + +function Component(props) { + const w = f(props.x); + const [x, _] = useState(w); + return
{x}
; +} + +function f(x) { + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.expect.md new file mode 100644 index 00000000000..dc97a4fab93 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.expect.md @@ -0,0 +1,78 @@ + +## Input + +```javascript +import { useState } from "react"; + +function Component(props) { + const w = f(props.x); + const [x, _] = useState(w); + return ( +
+ {x} + {w} +
+ ); +} + +function f(x) { + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; +import { useState } from "react"; + +function Component(props) { + const $ = _c(5); + let t0; + if ($[0] !== props.x) { + t0 = f(props.x); + $[0] = props.x; + $[1] = t0; + } else { + t0 = $[1]; + } + const w = t0; + const [x] = useState(w); + let t1; + if ($[2] !== x || $[3] !== w) { + t1 = ( +
+ {x} + {w} +
+ ); + $[2] = x; + $[3] = w; + $[4] = t1; + } else { + t1 = $[4]; + } + return t1; +} + +function f(x) { + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; + +``` + +### Eval output +(kind: ok)
4242
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.js new file mode 100644 index 00000000000..04bcb5cf6aa --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-unpruned-dependency.js @@ -0,0 +1,22 @@ +import { useState } from "react"; + +function Component(props) { + const w = f(props.x); + const [x, _] = useState(w); + return ( +
+ {x} + {w} +
+ ); +} + +function f(x) { + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ x: 42 }], + isComponent: true, +}; From e296fda158d7f22b813137e2eb25b77077f8567e Mon Sep 17 00:00:00 2001 From: Mike Vitousek Date: Wed, 29 May 2024 17:01:43 -0700 Subject: [PATCH 2/2] Update on "[compiler] Prune dependencies that are only used by useRef or useState" Summary: jmbrown215 recently had an observation that the arguments to useState/useRef are only used when a component renders for the first time, and never afterwards. We can skip more computation that we previously could, with reactive blocks that previously recomputed values when inputs changed now only ever computing them on the first render. [ghstack-poisoned]