From 2e3aa3954ce69db9b1e488d09cd49dcb2e5b7bf7 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Tue, 4 Apr 2023 12:30:10 -0700 Subject: [PATCH] Memoize hook args (treat as escaping) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Updates `PruneNonEscapingScopes` to consider hook arguments as potentially escaping. This is because hook inputs are "owned" by React — for example, closures passed to `useEffect`, or a value that is passed to a custom hook and which then becomes a memoized input. --- compiler/forget/src/HIR/Globals.ts | 18 +++++ .../ReactiveScopes/PruneNonEscapingScopes.ts | 29 +++++-- .../compiler/useEffect-arg-memoized.expect.md | 75 +++++++++++++++++++ .../compiler/useEffect-arg-memoized.js | 16 ++++ 4 files changed, 131 insertions(+), 7 deletions(-) create mode 100644 compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.expect.md create mode 100644 compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.js diff --git a/compiler/forget/src/HIR/Globals.ts b/compiler/forget/src/HIR/Globals.ts index e2f28c8f77..8770cc0ecb 100644 --- a/compiler/forget/src/HIR/Globals.ts +++ b/compiler/forget/src/HIR/Globals.ts @@ -169,6 +169,24 @@ const BUILTIN_HOOKS: Array<[string, Hook]> = [ valueKind: ValueKind.Frozen, }, ], + [ + "useEffect", + { + kind: "Memo", + name: "useEffect", + effectKind: Effect.Freeze, + valueKind: ValueKind.Frozen, + }, + ], + [ + "useLayoutEffect", + { + kind: "Memo", + name: "useLayoutEffect", + effectKind: Effect.Freeze, + valueKind: ValueKind.Frozen, + }, + ], ]; export type Global = BuiltInType | HookType | PolyType; diff --git a/compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts index 69bb34cb5d..263d03cace 100644 --- a/compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts +++ b/compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts @@ -12,6 +12,7 @@ import { Effect, IdentifierId, InstructionId, + isHookType, Pattern, Place, ReactiveFunction, @@ -23,6 +24,7 @@ import { ReactiveValue, ScopeId, } from "../HIR"; +import { eachInstructionValueOperand } from "../HIR/visitors"; import { log } from "../Utils/logger"; import { assertExhaustive } from "../Utils/utils"; import { getPlaceScope } from "./BuildReactiveBlocks"; @@ -37,8 +39,13 @@ import { /** * This pass prunes reactive scopes that are not necessary to bound downstream computation. - * Specifically, the pass identifies the set of identifiers which are directly returned by - * the function and/or transitively aliased by a return value - ie, values that "escape". + * Specifically, the pass identifies the set of identifiers which may "escape". Values can + * escape in one of two ways: + * * They are directly returned by the function and/or transitively aliased by a return + * value. + * * They are passed as input to a hook. This is because any value passed to a hook may + * have its referenced ultimately stored by React (ie, be aliased by an external value). + * For example, the closure passed to useEffect escapes. * * Example to build intuition: * @@ -96,7 +103,8 @@ import { * b. Conditional and logical expressions (and a few others) are conditinally aliased, * depending on whether their result value is aliased. * c. JSX is always unaliased (though its props children may be) - * 2. The same pass which builds the graph also stores the set of returned identifiers. + * 2. The same pass which builds the graph also stores the set of returned identifiers and set of + * identifiers passed as arguments to hooks. * 3. We traverse the graph starting from the returned identifiers and mark reachable dependencies * as escaping, based on the combination of the parent node's type and its children (eg a * conditional node with an aliased dep promotes to aliased). @@ -198,7 +206,7 @@ class State { identifiers: Map = new Map(); scopes: Map = new Map(); - returned: Set = new Set(); + escapingValues: Set = new Set(); /** * Declare a new identifier, used for function id and params @@ -302,8 +310,8 @@ function computeMemoizedIdentifiers(state: State): Set { } // Walk from the "roots" aka returned identifiers. - for (const returned of state.returned) { - visit(returned); + for (const value of state.escapingValues) { + visit(value); } return memoized; @@ -649,6 +657,13 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor { instruction.lvalue.identifier.id, instruction.value.place.identifier.id ); + } else if (instruction.value.kind === "CallExpression") { + const callee = instruction.value.callee; + if (isHookType(callee.identifier)) { + for (const operand of eachInstructionValueOperand(instruction.value)) { + state.escapingValues.add(operand.identifier.id); + } + } } } @@ -659,7 +674,7 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor { this.traverseTerminal(stmt, state); if (stmt.terminal.kind === "return" && stmt.terminal.value !== null) { - state.returned.add(stmt.terminal.value.identifier.id); + state.escapingValues.add(stmt.terminal.value.identifier.id); } } } diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.expect.md b/compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.expect.md new file mode 100644 index 0000000000..2a8c6040f7 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.expect.md @@ -0,0 +1,75 @@ + +## Input + +```javascript +function Component(props) { + const dispatch = useDispatch(); + useFreeze(dispatch); + + // onUpdate should be memoized even though it doesn't + // flow into the return value + const onUpdate = () => { + dispatch({ kind: "update" }); + }; + + useEffect(() => { + onUpdate(); + }, [onUpdate]); + + return
; +} + +``` + +## Code + +```javascript +function Component(props) { + const $ = React.unstable_useMemoCache(7); + const dispatch = useDispatch(); + useFreeze(dispatch); + const c_0 = $[0] !== dispatch; + let t0; + if (c_0) { + t0 = () => { + dispatch({ kind: "update" }); + }; + $[0] = dispatch; + $[1] = t0; + } else { + t0 = $[1]; + } + const onUpdate = t0; + const c_2 = $[2] !== onUpdate; + let t1; + if (c_2) { + t1 = () => { + onUpdate(); + }; + $[2] = onUpdate; + $[3] = t1; + } else { + t1 = $[3]; + } + const c_4 = $[4] !== onUpdate; + let t2; + if (c_4) { + t2 = [onUpdate]; + $[4] = onUpdate; + $[5] = t2; + } else { + t2 = $[5]; + } + useEffect(t1, t2); + let t3; + if ($[6] === Symbol.for("react.memo_cache_sentinel")) { + t3 =
; + $[6] = t3; + } else { + t3 = $[6]; + } + return t3; +} + +``` + \ No newline at end of file diff --git a/compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.js b/compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.js new file mode 100644 index 0000000000..5af08ec407 --- /dev/null +++ b/compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.js @@ -0,0 +1,16 @@ +function Component(props) { + const dispatch = useDispatch(); + useFreeze(dispatch); + + // onUpdate should be memoized even though it doesn't + // flow into the return value + const onUpdate = () => { + dispatch({ kind: "update" }); + }; + + useEffect(() => { + onUpdate(); + }, [onUpdate]); + + return
; +}