From f7d16db54480558fe792d04d44ae3b8d18603894 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Wed, 29 Nov 2023 10:46:43 -0800 Subject: [PATCH] [RFC] Refine memoization for Array#map with non-mutating callbacks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Improves memoization for cases such as #2409: ```javascript const x = []; useEffect(...); return
{x.map(item => {item})}
; ``` We previously thought that the `x.map(...)` call mutated `x` since its kind was Mutable. However, in this case we can determine that the map call cannot mutate `x` (or anything else): the lambda does not mutate any free variables and does not mutate its arguments. This PR adds a new flag to function signatures, used for method calls only, that checks for such cases. The idea is that if the receiver is the only thing that is mutable — including that there are no args which are function expressions which mutate their parameters — then we can infer the effect as a read. See tests which confirm that function expressions which capture or mutate their params bypass the optimization. --- .../babel-plugin-react-forget/src/HIR/HIR.ts | 1 + .../src/HIR/ObjectShape.ts | 79 ++++++++++++++ .../src/Inference/InferReferenceEffects.ts | 76 +++++++++++++ ...ay-map-captures-receiver-noAlias.expect.md | 26 ++++- ...y-with-capturing-map-after-hook.expect.md} | 31 ++++-- ...ize-array-with-capturing-map-after-hook.js | 26 +++++ ...rray-with-mutable-map-after-hook.expect.md | 102 ++++++++++++++++++ ...moize-array-with-mutable-map-after-hook.js | 26 +++++ ...ay-with-immutable-map-after-hook.expect.md | 102 ++++++++++++++++++ ...ze-array-with-immutable-map-after-hook.js} | 0 10 files changed, 456 insertions(+), 13 deletions(-) rename compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/{repro-missing-memoization-unmodified-array.expect.md => repro-dont-memoize-array-with-capturing-map-after-hook.expect.md} (63%) create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-capturing-map-after-hook.js create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.expect.md create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.js create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-memoize-array-with-immutable-map-after-hook.expect.md rename compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/{repro-missing-memoization-unmodified-array.js => repro-memoize-array-with-immutable-map-after-hook.js} (100%) diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts index 8e5b08e28e..dfb022e733 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts @@ -971,6 +971,7 @@ export enum Effect { * But we do not error if the value is known to be immutable. */ ConditionallyMutate = "mutate?", + /* * This reference *does* write to (mutate) the value. It is an error (invalid input) * if an immutable value flows into a location with this effect. 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 8467a76d70..5910f7327a 100644 --- a/compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts +++ b/compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts @@ -147,6 +147,20 @@ export type FunctionSignature = { * may choose not to memoize arguments if they do not otherwise escape. */ noAlias?: boolean; + + /** + * Supported only for methods (no-op when used on functions in CallExpression.callee position). + * + * Indicates that the method can only modify its receiver if any of the arguments + * are mutable or are function expressions which mutate their arguments. This is designed + * for methods such as Array.prototype.map(), which only mutate the receiver array if they are + * passed a callback which has mutable side-effects (including mutating its inputs). + * + * MethodCalls to such functions will use a different behavior depending on their arguments: + * - If arguments are all non-mutable, the arguments get the Read effect and the receiver is Capture. + * - Else uses the effects specified by this signature. + */ + mutableOnlyIfOperandsAreMutable?: boolean; }; /* @@ -231,6 +245,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [ calleeEffect: Effect.ConditionallyMutate, returnValueKind: ValueKind.Mutable, noAlias: true, + mutableOnlyIfOperandsAreMutable: true, }), ], [ @@ -247,6 +262,70 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [ calleeEffect: Effect.ConditionallyMutate, returnValueKind: ValueKind.Mutable, noAlias: true, + mutableOnlyIfOperandsAreMutable: true, + }), + ], + [ + "every", + addFunction(BUILTIN_SHAPES, [], { + positionalParams: [], + restParam: Effect.ConditionallyMutate, + returnType: { kind: "Primitive" }, + /* + * callee is ConditionallyMutate because items of the array + * flow into the lambda and may be mutated there, even though + * the array object itself is not modified + */ + calleeEffect: Effect.ConditionallyMutate, + returnValueKind: ValueKind.Immutable, + noAlias: true, + mutableOnlyIfOperandsAreMutable: true, + }), + ], + [ + "some", + addFunction(BUILTIN_SHAPES, [], { + positionalParams: [], + restParam: Effect.ConditionallyMutate, + returnType: { kind: "Primitive" }, + /* + * callee is ConditionallyMutate because items of the array + * flow into the lambda and may be mutated there, even though + * the array object itself is not modified + */ + calleeEffect: Effect.ConditionallyMutate, + returnValueKind: ValueKind.Immutable, + noAlias: true, + mutableOnlyIfOperandsAreMutable: true, + }), + ], + [ + "find", + addFunction(BUILTIN_SHAPES, [], { + positionalParams: [], + restParam: Effect.ConditionallyMutate, + returnType: { kind: "Poly" }, + calleeEffect: Effect.ConditionallyMutate, + returnValueKind: ValueKind.Mutable, + noAlias: true, + mutableOnlyIfOperandsAreMutable: true, + }), + ], + [ + "findIndex", + addFunction(BUILTIN_SHAPES, [], { + positionalParams: [], + restParam: Effect.ConditionallyMutate, + returnType: { kind: "Primitive" }, + /* + * callee is ConditionallyMutate because items of the array + * flow into the lambda and may be mutated there, even though + * the array object itself is not modified + */ + calleeEffect: Effect.ConditionallyMutate, + returnValueKind: ValueKind.Immutable, + noAlias: true, + mutableOnlyIfOperandsAreMutable: true, }), ], [ diff --git a/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts b/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts index a61219fae4..278ab358ed 100644 --- a/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts +++ b/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts @@ -228,6 +228,17 @@ class InferenceState { this.#values.set(value, kind); } + values(place: Place): Array { + const values = this.#variables.get(place.identifier.id); + CompilerError.invariant(values != null, { + reason: `[hoisting] Expected value kind to be initialized`, + description: `${printPlace(place)}`, + loc: place.loc, + suggestions: null, + }); + return Array.from(values); + } + // Lookup the kind of the given @param value. kind(place: Place): ValueKind { const values = this.#variables.get(place.identifier.id); @@ -833,6 +844,26 @@ function inferBlock( instrValue.property.identifier.type ); + if ( + signature !== null && + signature.mutableOnlyIfOperandsAreMutable && + areArgumentsImmutableAndNonMutating(state, instrValue.args) + ) { + /* + * None of the args are mutable or mutate their params, we can downgrade to + * treating as all reads + */ + for (const arg of instrValue.args) { + const place = arg.kind === "Identifier" ? arg : arg.place; + state.reference(place, Effect.Read); + } + state.reference(instrValue.receiver, Effect.Read); + state.initialize(instrValue, signature.returnValueKind); + state.define(instr.lvalue, instrValue); + instr.lvalue.effect = Effect.ConditionallyMutate; + continue; + } + const effects = signature !== null ? getFunctionEffects(instrValue, signature) : null; const returnValueKind = @@ -1185,3 +1216,48 @@ function getFunctionEffects( } return results; } + +/** + * Returns true if all of the arguments are both non-mutable (immutable or frozen) + * _and_ are not functions which might mutate their arguments. Note that function + * expressions count as frozen so long as they do not mutate free variables: this + * function checks that such functions also don't mutate their inputs. + */ +function areArgumentsImmutableAndNonMutating( + state: InferenceState, + args: MethodCall["args"] +): boolean { + for (const arg of args) { + const place = arg.kind === "Identifier" ? arg : arg.place; + const kind = state.kind(place); + switch (kind) { + case ValueKind.Immutable: + case ValueKind.Frozen: { + /* + * Only immutable values, or frozen lambdas are allowed. + * A lambda may appear frozen even if it may mutate its inputs, + * so we have a second check even for frozen value types + */ + break; + } + default: { + return false; + } + } + const values = state.values(place); + for (const value of values) { + if ( + value.kind === "FunctionExpression" && + value.loweredFunc.func.params.some((param) => { + const place = param.kind === "Identifier" ? param : param.place; + const range = place.identifier.mutableRange; + return range.end > range.start + 1; + }) + ) { + // This is a function which may mutate its inputs + return false; + } + } + } + return true; +} diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-captures-receiver-noAlias.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-captures-receiver-noAlias.expect.md index e85e2af46b..80dfe79364 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-captures-receiver-noAlias.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-captures-receiver-noAlias.expect.md @@ -23,18 +23,34 @@ export const FIXTURE_ENTRYPOINT = { ```javascript import { unstable_useMemoCache as useMemoCache } from "react"; function Component(props) { - const $ = useMemoCache(2); + const $ = useMemoCache(6); let t0; if ($[0] !== props.a) { - const item = { a: props.a }; - const items = [item]; - t0 = items.map((item_0) => item_0); + t0 = { a: props.a }; $[0] = props.a; $[1] = t0; } else { t0 = $[1]; } - const mapped = t0; + const item = t0; + let t1; + if ($[2] !== item) { + t1 = [item]; + $[2] = item; + $[3] = t1; + } else { + t1 = $[3]; + } + const items = t1; + let t2; + if ($[4] !== items) { + t2 = items.map((item_0) => item_0); + $[4] = items; + $[5] = t2; + } else { + t2 = $[5]; + } + const mapped = t2; return mapped; } diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-missing-memoization-unmodified-array.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-capturing-map-after-hook.expect.md similarity index 63% rename from compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-missing-memoization-unmodified-array.expect.md rename to compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-capturing-map-after-hook.expect.md index 9cd0e01721..b0f65cbb3d 100644 --- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-missing-memoization-unmodified-array.expect.md +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-capturing-map-after-hook.expect.md @@ -3,25 +3,29 @@ ```javascript import { useEffect, useState } from "react"; +import { mutate } from "shared-runtime"; function Component(props) { - const x = [props.value]; + const x = [{ ...props.value }]; useEffect(() => {}, []); const onClick = () => { console.log(x.length); }; + let y; return (
{x.map((item) => { - return {item}; + y = item; + return {item.text}; })} + {mutate(y)}
); } export const FIXTURE_ENTRYPOINT = { fn: Component, - params: [{ value: 42 }], + params: [{ value: { id: 0, text: "Hello!" } }], isComponent: true, }; @@ -35,10 +39,11 @@ import { useState, unstable_useMemoCache as useMemoCache, } from "react"; +import { mutate } from "shared-runtime"; function Component(props) { const $ = useMemoCache(5); - const x = [props.value]; + const x = [{ ...props.value }]; let t0; let t1; if ($[0] === Symbol.for("react.memo_cache_sentinel")) { @@ -55,10 +60,20 @@ function Component(props) { console.log(x.length); }; - const t2 = x.map((item) => {item}); + let y; + + const t2 = x.map((item) => { + y = item; + return {item.text}; + }); let t3; if ($[2] !== onClick || $[3] !== t2) { - t3 =
{t2}
; + t3 = ( +
+ {t2} + {mutate(y)} +
+ ); $[2] = onClick; $[3] = t2; $[4] = t3; @@ -70,11 +85,11 @@ function Component(props) { export const FIXTURE_ENTRYPOINT = { fn: Component, - params: [{ value: 42 }], + params: [{ value: { id: 0, text: "Hello!" } }], isComponent: true, }; ``` ### Eval output -(kind: ok)
42
\ No newline at end of file +(kind: ok)
Hello!
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-capturing-map-after-hook.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-capturing-map-after-hook.js new file mode 100644 index 0000000000..ae1ebf63b8 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-capturing-map-after-hook.js @@ -0,0 +1,26 @@ +import { useEffect, useState } from "react"; +import { mutate } from "shared-runtime"; + +function Component(props) { + const x = [{ ...props.value }]; + useEffect(() => {}, []); + const onClick = () => { + console.log(x.length); + }; + let y; + return ( +
+ {x.map((item) => { + y = item; + return {item.text}; + })} + {mutate(y)} +
+ ); +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: { id: 0, text: "Hello!" } }], + isComponent: true, +}; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.expect.md new file mode 100644 index 0000000000..c869711dae --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.expect.md @@ -0,0 +1,102 @@ + +## Input + +```javascript +import { useEffect, useState } from "react"; +import { mutate } from "shared-runtime"; + +function Component(props) { + const x = [{ ...props.value }]; + useEffect(() => {}, []); + const onClick = () => { + console.log(x.length); + }; + let y; + return ( +
+ {x.map((item) => { + item.flag = true; + return {item.text}; + })} + {mutate(y)} +
+ ); +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: { id: 0, text: "Hello", flag: false } }], + isComponent: true, +}; + +``` + +## Code + +```javascript +import { + useEffect, + useState, + unstable_useMemoCache as useMemoCache, +} from "react"; +import { mutate } from "shared-runtime"; + +function Component(props) { + const $ = useMemoCache(6); + const x = [{ ...props.value }]; + let t0; + let t1; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t0 = () => {}; + t1 = []; + $[0] = t0; + $[1] = t1; + } else { + t0 = $[0]; + t1 = $[1]; + } + useEffect(t0, t1); + const onClick = () => { + console.log(x.length); + }; + + let y; + + const t3 = x.map((item) => { + item.flag = true; + return {item.text}; + }); + let t2; + if ($[2] === Symbol.for("react.memo_cache_sentinel")) { + t2 = mutate(y); + $[2] = t2; + } else { + t2 = $[2]; + } + let t4; + if ($[3] !== onClick || $[4] !== t3) { + t4 = ( +
+ {t3} + {t2} +
+ ); + $[3] = onClick; + $[4] = t3; + $[5] = t4; + } else { + t4 = $[5]; + } + return t4; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: { id: 0, text: "Hello", flag: false } }], + isComponent: true, +}; + +``` + +### Eval output +(kind: ok)
Hello
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.js new file mode 100644 index 0000000000..76b2691ad0 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.js @@ -0,0 +1,26 @@ +import { useEffect, useState } from "react"; +import { mutate } from "shared-runtime"; + +function Component(props) { + const x = [{ ...props.value }]; + useEffect(() => {}, []); + const onClick = () => { + console.log(x.length); + }; + let y; + return ( +
+ {x.map((item) => { + item.flag = true; + return {item.text}; + })} + {mutate(y)} +
+ ); +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: { id: 0, text: "Hello", flag: false } }], + isComponent: true, +}; diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-memoize-array-with-immutable-map-after-hook.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-memoize-array-with-immutable-map-after-hook.expect.md new file mode 100644 index 0000000000..a78d133011 --- /dev/null +++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-memoize-array-with-immutable-map-after-hook.expect.md @@ -0,0 +1,102 @@ + +## Input + +```javascript +import { useEffect, useState } from "react"; + +function Component(props) { + const x = [props.value]; + useEffect(() => {}, []); + const onClick = () => { + console.log(x.length); + }; + return ( +
+ {x.map((item) => { + return {item}; + })} +
+ ); +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], + isComponent: true, +}; + +``` + +## Code + +```javascript +import { + useEffect, + useState, + unstable_useMemoCache as useMemoCache, +} from "react"; + +function Component(props) { + const $ = useMemoCache(11); + let t0; + if ($[0] !== props.value) { + t0 = [props.value]; + $[0] = props.value; + $[1] = t0; + } else { + t0 = $[1]; + } + const x = t0; + let t1; + let t2; + if ($[2] === Symbol.for("react.memo_cache_sentinel")) { + t1 = () => {}; + t2 = []; + $[2] = t1; + $[3] = t2; + } else { + t1 = $[2]; + t2 = $[3]; + } + useEffect(t1, t2); + let t3; + if ($[4] !== x.length) { + t3 = () => { + console.log(x.length); + }; + $[4] = x.length; + $[5] = t3; + } else { + t3 = $[5]; + } + const onClick = t3; + let t4; + if ($[6] !== x) { + t4 = x.map((item) => {item}); + $[6] = x; + $[7] = t4; + } else { + t4 = $[7]; + } + let t5; + if ($[8] !== onClick || $[9] !== t4) { + t5 =
{t4}
; + $[8] = onClick; + $[9] = t4; + $[10] = t5; + } else { + t5 = $[10]; + } + return t5; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ value: 42 }], + isComponent: true, +}; + +``` + +### Eval output +(kind: ok)
42
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-missing-memoization-unmodified-array.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-memoize-array-with-immutable-map-after-hook.js similarity index 100% rename from compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-missing-memoization-unmodified-array.js rename to compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-memoize-array-with-immutable-map-after-hook.js