From 1879fbe644105de155eda887bfae733a5f5a80ca Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Fri, 15 Sep 2023 16:18:03 -0700 Subject: [PATCH] Option to not memoize non-aliased function parameters This has been nagging at me for a _long_ time: we unnecessarily memoize function callbacks passed to things like Array.prototype.map, even though we know these functions can't escape. This PR fixes this as follows: * Adds a `noAlias?: boolean` flag to builtin function signatures, defaulting to false if not specified. * Adds a feature flag, `enableNoAliasOptimizations`, to gate optimizations based on the value of that new flag. * When the feature is enabled, `PruneNonEscapingScopes` now looks up the signature of method calls, and avoids memoizing the arguments if the signatures specifies `noAlias: true`. * Annotates Array.prototype.map and Array.prototype.filter as `noAlias`. This does not mean we'll never memoize arguments to Array.prototype.map, it just means that the argument itself won't be considered as escaping. If the function still escapes by some other means it will get memoized: ``` function Component(props) { const f = () => {}; // memoized! const x = [].map(f); // not from here.. return [x, f]; // but bc it escapes here } ``` Note: this delivers some of the wins from #1640. That PR tried to do a bunch of things, part of which I already landed w the introduction of ConditionallyMutate, which allowed us to type Array.prototype.map. This PR further gives us the ability to understand functions that don't alias their params at all. The remaining bit from #1640 is the idea of understanding that key hooks such as `useFragment()` return transitively readonly, transitively array/object/primitive values, and any `.map()` or `.filter()` calls must be on arrays, allowing us to optimize them. Without that extra step, we'll still have to memoize a lot of `array.map()` lambdas just because we aren't sure that the receiver is an Array. But this PR helps with some cases, and lays the groundwork for the rest of that PR. --- .../src/HIR/Environment.ts | 12 ++++ .../src/HIR/ObjectShape.ts | 8 +++ .../ReactiveScopes/PruneNonEscapingScopes.ts | 56 ++++++++++++++--- .../array-map-frozen-array-noAlias.expect.md | 60 ++++++++++++++++++ .../array-map-frozen-array-noAlias.js | 13 ++++ ...le-array-mutating-lambda-noAlias.expect.md | 61 ++++++++++++++++++ ...p-mutable-array-mutating-lambda-noAlias.js | 15 +++++ ...ay-map-noAlias-escaping-function.expect.md | 62 +++++++++++++++++++ .../array-map-noAlias-escaping-function.js | 11 ++++ .../fixture-test-utils/src/compiler-utils.ts | 7 ++- 10 files changed, 294 insertions(+), 11 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.expect.md create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.js create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.expect.md create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.js create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.expect.md create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.js diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts index ffd4d43ae7..20a8e2cdf6 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts @@ -109,6 +109,15 @@ export type EnvironmentConfig = Partial<{ */ enableFunctionCallSignatureOptimizations: boolean; + /** + * Enable optimizations based on the `noAlias` flag of method signatures. When enabled, + * function signatures can declare that they do not alias their arguments, allowing + * Forget to (in some cases) avoid memoizing arguments if they do not otherwise escape. + * + * Defaults to false + */ + enableNoAliasOptimizations: boolean; + /** * When enabled, the compiler assumes that hooks follow the Rules of React: * - Hooks may memoize computation based on any of their parameters, thus @@ -182,6 +191,7 @@ export class Environment { enableFunctionCallSignatureOptimizations: boolean; enableAssumeHooksFollowRulesOfReact: boolean; enableTreatHooksAsFunctions: boolean; + enableNoAliasOptimizations: boolean; disableAllMemoization: boolean; enableEmitFreeze: ExternalFunction | null; assertValidMutableRanges: boolean; @@ -225,6 +235,8 @@ export class Environment { this.validateFrozenLambdas = config?.validateFrozenLambdas ?? false; this.enableFunctionCallSignatureOptimizations = config?.enableFunctionCallSignatureOptimizations ?? false; + this.enableNoAliasOptimizations = + config?.enableNoAliasOptimizations ?? false; this.enableAssumeHooksFollowRulesOfReact = config?.enableAssumeHooksFollowRulesOfReact ?? false; this.enableTreatHooksAsFunctions = diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts index 34262d7aed..66e1c94845 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts @@ -138,6 +138,12 @@ export type FunctionSignature = { returnValueKind: ValueKind; calleeEffect: Effect; hookKind: HookKind | null; + /** + * Whether any of the parameters may be aliased by each other or the return + * value. Defaults to false (parameters may alias). When true, the compiler + * may choose not to memoize arguments if they do not otherwise escape. + */ + noAlias?: boolean; }; /** @@ -217,6 +223,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [ // the array object itself is not modified calleeEffect: Effect.ConditionallyMutate, returnValueKind: ValueKind.Mutable, + noAlias: true, }), ], [ @@ -230,6 +237,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [ // the array object itself is not modified calleeEffect: Effect.ConditionallyMutate, returnValueKind: ValueKind.Mutable, + noAlias: true, }), ], [ diff --git a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts index 4652d9db10..f0ab3e20e9 100644 --- a/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts +++ b/compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts @@ -24,6 +24,7 @@ import { getHookKind, isMutableEffect, } from "../HIR"; +import { getFunctionCallSignature } from "../Inference/InferReferenceEffects"; import { log } from "../Utils/logger"; import { assertExhaustive } from "../Utils/utils"; import { getPlaceScope } from "./BuildReactiveBlocks"; @@ -119,7 +120,11 @@ export function pruneNonEscapingScopes( for (const param of fn.params) { state.declare(param.identifier.id); } - visitReactiveFunction(fn, new CollectDependenciesVisitor(options), state); + visitReactiveFunction( + fn, + new CollectDependenciesVisitor(fn.env, options), + state + ); log(() => prettyFormat(state)); @@ -343,6 +348,7 @@ type LValueMemoization = { * - level: the level of memoization to apply to this value */ function computeMemoizationInputs( + env: Environment, value: ReactiveValue, lvalue: Place | null, options: MemoizationOptions @@ -361,8 +367,10 @@ function computeMemoizationInputs( : [], rvalues: [ // Conditionals do not alias their test value. - ...computeMemoizationInputs(value.consequent, null, options).rvalues, - ...computeMemoizationInputs(value.alternate, null, options).rvalues, + ...computeMemoizationInputs(env, value.consequent, null, options) + .rvalues, + ...computeMemoizationInputs(env, value.alternate, null, options) + .rvalues, ], }; } @@ -374,8 +382,8 @@ function computeMemoizationInputs( ? [{ place: lvalue, level: MemoizationLevel.Conditional }] : [], rvalues: [ - ...computeMemoizationInputs(value.left, null, options).rvalues, - ...computeMemoizationInputs(value.right, null, options).rvalues, + ...computeMemoizationInputs(env, value.left, null, options).rvalues, + ...computeMemoizationInputs(env, value.right, null, options).rvalues, ], }; } @@ -389,7 +397,8 @@ function computeMemoizationInputs( // Only the final value of the sequence is a true rvalue: // values from the sequence's instructions are evaluated // as separate nodes - rvalues: computeMemoizationInputs(value.value, null, options).rvalues, + rvalues: computeMemoizationInputs(env, value.value, null, options) + .rvalues, }; } case "JsxExpression": { @@ -595,18 +604,42 @@ function computeMemoizationInputs( return { lvalues: lvalues, rvalues: [ - ...computeMemoizationInputs(value.value, null, options).rvalues, + ...computeMemoizationInputs(env, value.value, null, options).rvalues, ], }; } + case "MethodCall": { + const signature = env.enableNoAliasOptimizations + ? getFunctionCallSignature(env, value.property.identifier.type) + : null; + const operands = [...eachReactiveValueOperand(value)]; + let lvalues = []; + if (lvalue !== null) { + lvalues.push({ place: lvalue, level: MemoizationLevel.Memoized }); + } + if (signature?.noAlias === true) { + return { + lvalues, + rvalues: [], + }; + } + lvalues.push( + ...operands + .filter((operand) => isMutableEffect(operand.effect, operand.loc)) + .map((place) => ({ place, level: MemoizationLevel.Memoized })) + ); + return { + lvalues, + rvalues: operands, + }; + } case "RegExpLiteral": case "FunctionExpression": case "TaggedTemplateExpression": - case "CallExpression": case "ArrayExpression": case "NewExpression": case "ObjectExpression": - case "MethodCall": + case "CallExpression": case "PropertyStore": { // All of these instructions may produce new values which must be memoized if // reachable from a return value. Any mutable rvalue may alias any other rvalue @@ -680,10 +713,12 @@ function computePatternLValues(pattern: Pattern): Array { * identifier's and scope's dependencies. */ class CollectDependenciesVisitor extends ReactiveFunctionVisitor { + env: Environment; options: MemoizationOptions; - constructor(options: MemoizationOptions) { + constructor(env: Environment, options: MemoizationOptions) { super(); + this.env = env; this.options = options; } @@ -695,6 +730,7 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor { // Determe the level of memoization for this value and the lvalues/rvalues const aliasing = computeMemoizationInputs( + this.env, instruction.value, instruction.lvalue, this.options diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.expect.md new file mode 100644 index 0000000000..e14a823f42 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.expect.md @@ -0,0 +1,60 @@ + +## Input + +```javascript +// @enableNoAliasOptimizations +function Component(props) { + const x = []; + {x}; + const y = x.map((item) => item); + return [x, y]; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], + isComponent: false, +}; + +``` + +## Code + +```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; // @enableNoAliasOptimizations +function Component(props) { + const $ = useMemoCache(3); + let t0; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t0 = []; + $[0] = t0; + } else { + t0 = $[0]; + } + const x = t0; + let t1; + if ($[1] === Symbol.for("react.memo_cache_sentinel")) { + t1 = x.map((item) => item); + $[1] = t1; + } else { + t1 = $[1]; + } + const y = t1; + let t2; + if ($[2] === Symbol.for("react.memo_cache_sentinel")) { + t2 = [x, y]; + $[2] = t2; + } else { + t2 = $[2]; + } + return t2; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], + isComponent: false, +}; + +``` + \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.js new file mode 100644 index 0000000000..0d8df82fd2 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.js @@ -0,0 +1,13 @@ +// @enableNoAliasOptimizations +function Component(props) { + const x = []; + {x}; + const y = x.map((item) => item); + return [x, y]; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], + isComponent: false, +}; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.expect.md new file mode 100644 index 0000000000..62ef2ade57 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.expect.md @@ -0,0 +1,61 @@ + +## Input + +```javascript +// @enableNoAliasOptimizations +function Component(props) { + const x = []; + const y = x.map((item) => { + item.updated = true; + return item; + }); + return [x, y]; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], + isComponent: false, +}; + +``` + +## Code + +```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; // @enableNoAliasOptimizations +function Component(props) { + const $ = useMemoCache(3); + let t0; + let x; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + x = []; + t0 = x.map((item) => { + item.updated = true; + return item; + }); + $[0] = t0; + $[1] = x; + } else { + t0 = $[0]; + x = $[1]; + } + const y = t0; + let t1; + if ($[2] === Symbol.for("react.memo_cache_sentinel")) { + t1 = [x, y]; + $[2] = t1; + } else { + t1 = $[2]; + } + return t1; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], + isComponent: false, +}; + +``` + \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.js new file mode 100644 index 0000000000..68b90eda4e --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.js @@ -0,0 +1,15 @@ +// @enableNoAliasOptimizations +function Component(props) { + const x = []; + const y = x.map((item) => { + item.updated = true; + return item; + }); + return [x, y]; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{}], + isComponent: false, +}; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.expect.md new file mode 100644 index 0000000000..af22a304df --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.expect.md @@ -0,0 +1,62 @@ + +## Input + +```javascript +function Component(props) { + const f = (item) => item; + const x = [...props.items].map(f); // `f` doesn't escape here... + return [x, f]; // ...but it does here so it's memoized +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ items: [{ id: 1 }] }], + isComponent: false, +}; + +``` + +## Code + +```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; +function Component(props) { + const $ = useMemoCache(5); + let t0; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t0 = (item) => item; + $[0] = t0; + } else { + t0 = $[0]; + } + const f = t0; + const c_1 = $[1] !== props.items; + let t1; + if (c_1) { + t1 = [...props.items].map(f); + $[1] = props.items; + $[2] = t1; + } else { + t1 = $[2]; + } + const x = t1; + const c_3 = $[3] !== x; + let t2; + if (c_3) { + t2 = [x, f]; + $[3] = x; + $[4] = t2; + } else { + t2 = $[4]; + } + return t2; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ items: [{ id: 1 }] }], + isComponent: false, +}; + +``` + \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.js new file mode 100644 index 0000000000..60726b5e5c --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.js @@ -0,0 +1,11 @@ +function Component(props) { + const f = (item) => item; + const x = [...props.items].map(f); // `f` doesn't escape here... + return [x, f]; // ...but it does here so it's memoized +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ items: [{ id: 1 }] }], + isComponent: false, +}; diff --git a/compiler/packages/fixture-test-utils/src/compiler-utils.ts b/compiler/packages/fixture-test-utils/src/compiler-utils.ts index 44e5c045e2..7582a419df 100644 --- a/compiler/packages/fixture-test-utils/src/compiler-utils.ts +++ b/compiler/packages/fixture-test-utils/src/compiler-utils.ts @@ -28,6 +28,7 @@ export function transformFixtureInput( let enableEmitFreeze = null; let compilationMode: CompilationMode = "all"; let enableForest = false; + let enableNoAliasOptimizations = false; if (firstLine.indexOf("@compilationMode(annotation)") !== -1) { assert( @@ -83,6 +84,9 @@ export function transformFixtureInput( if (firstLine.includes("@enableForest true")) { enableForest = true; } + if (firstLine.includes("@enableNoAliasOptimizations")) { + enableNoAliasOptimizations = true; + } return pluginFn( input, @@ -101,6 +105,7 @@ export function transformFixtureInput( ]), enableAssumeHooksFollowRulesOfReact, enableFunctionCallSignatureOptimizations: true, + enableNoAliasOptimizations, disableAllMemoization, enableTreatHooksAsFunctions, inlineUseMemo: true, @@ -117,7 +122,7 @@ export function transformFixtureInput( logger: null, gating, instrumentForget, - panicThreshold: 'ALL_ERRORS', + panicThreshold: "ALL_ERRORS", noEmit: false, }, includeAst