@samitouri / QOS-React-2 / commits / 2abd439b43

Option to preserve existing memoization guarantees

Adds an option to preserve existing memoization guarantees for values produced with useMemo and useCallback. We still discard the calls to these hooks, but we preserve the information that the value is frozen at that point in the program. Because these values are produced solely within the useMemo/useCallback callback, their mutation cannot have any interspersed hook calls. This means that the values mutable range will never span a hook and end at the point of the useMemo, ensuring that they are memoized at the same point. The main things that can change (relative to the orignal code) are: * Forget will infer a precise set of dependencies, ignoring the user-provided values. In practice this should only occur if the original code had a lint violation, which Forget would bail out on. So in practice this shouldn't happen unless the code doesn't use the React linter. * Forget may start the memoization block earlier than the developer did if other values are mutated along with the value being produced. This can cause memoization to fail, but only in situations where it would have failed previously: ```javascript const a = []; useFoo(); const b = useMemo(() => { const c = a; c.push(1); return c; }, [a]); ``` In this example (sans Forget) the useMemo will invalidate on every render because `a` will always be a new array and its listed as a dependency of the useMemo. Forget would correctly determine that the memoization would have to work as follows: ```javascript let c; if (...) { const a = [] useFoo(); // OOPS we made a hook call conditional const t0 = a; t0.push(1); c = t0; ... } else { c = $[...] } ``` Because this is invalid, Forget would (later in the pipeline) strip out this memoization block and (as with the original) leave `c` un-memoized. In this same example, removing the hook would cause Forget to be able to memoize a value that wasn't memoized before: ```javascript const a = []; const b = useMemo(() => { const c = a; c.push(1); return c; }, [a]); ``` This invalidates every render without Forget, but would memoize correctly with Forget (it would expand the memoization block to include the declaration of `a`).

Joe Savona committed Dec 15, 2023 at 13:47 UTC 2abd439b43679f77c2fb9ba2f38907fa37f953eb

16 files changed +337 -19

compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts

+20

@@ -102,6 +102,26 @@ export type Hook = z.infer<typeof HookSchema>;

102 const EnvironmentConfigSchema = z.object({

103 customHooks: z.map(z.string(), HookSchema).optional().default(new Map()),

104

105 + /**

106 + * Enable using information from existing useMemo/useCallback to understand when a value is done

107 + * being mutated. With this mode enabled, Forget will still discard the actual useMemo/useCallback

108 + * calls and may memoize slightly differently. However, it will assume that the values produced

109 + * are not subsequently modified, guaranteeing that the value will be memoized.

110 + *

111 + * By preserving guarantees about when values are memoized, this option preserves any existing

112 + * behavior that depends on referential equality in the original program. Notably, this preserves

113 + * existing effect behavior (how often effects fire) for effects that rely on referential equality.

114 + *

115 + * When disabled, Forget will not only prune useMemo and useCallback calls but also completely ignore

116 + * them, not using any information from them to guide compilation. Therefore, disabling this flag

117 + * will produce output that mimics the result from removing all memoization.

118 + *

119 + * Our recommendation is to first try running your application with this flag enabled, then attempt

120 + * to disable this flag and see what changes or breaks. This will mostly likely be effects that

121 + * depend on referential equality, which can be refactored (TODO guide for this).

122 + */

123 + enablePreserveExistingMemoizationGuarantees: z.boolean().default(false),

124 +

125 // 🌲

126 enableForest: z.boolean().default(false),

127 // <🌲>

compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts

@@ -845,6 +845,12 @@ export type InstructionValue =

845 }

846 // `debugger` statement

847 | { kind: "Debugger"; loc: SourceLocation }

848 + /*

849 + * Represents semantic information from useMemo/useCallback that the developer

850 + * has indicated a particular value should be memoized. This value is ignored

851 + * unless the TODO flag is enabled.

852 + */

853 + | { kind: "Memoize"; value: Place; loc: SourceLocation }

854 /*

855 * Catch-all for statements such as type imports, nested class declarations, etc

856 * which are not directly represented, but included for completeness and to allow

compiler/packages/babel-plugin-react-forget/src/HIR/HIRBuilder.ts

+19

@@ -15,11 +15,14 @@ import {

15 BasicBlock,

16 BlockId,

17 BlockKind,

18 + Effect,

19 + GeneratedSource,

20 GotoVariant,

21 HIR,

22 Identifier,

23 IdentifierId,

24 Instruction,

25 + Place,

26 Terminal,

27 makeBlockId,

28 makeInstructionId,

@@ -856,3 +859,19 @@ export function removeUnnecessaryTryCatch(fn: HIR): void {

859 }

860 }

861 }

862 +

863 +export function createTemporaryPlace(env: Environment): Place {

864 + return {

865 + kind: "Identifier",

866 + identifier: {

867 + id: env.nextIdentifierId,

868 + mutableRange: { start: makeInstructionId(0), end: makeInstructionId(0) },

869 + name: null,

870 + scope: null,

871 + type: makeType(),

872 + },

873 + reactive: false,

874 + effect: Effect.Unknown,

875 + loc: GeneratedSource,

876 + };

877 +}

compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts

@@ -594,6 +594,10 @@ export function printInstructionValue(instrValue: ReactiveValue): string {

594 } ${printPlace(instrValue.value)}`;

595 break;

596 }

597 + case "Memoize": {

598 + value = `Memoize ${printPlace(instrValue.value)}`;

599 + break;

600 + }

601 default: {

602 assertExhaustive(

603 instrValue,

compiler/packages/babel-plugin-react-forget/src/HIR/visitors.ts

@@ -209,6 +209,10 @@ export function* eachInstructionValueOperand(

209 yield instrValue.value;

210 break;

211 }

212 + case "Memoize": {

213 + yield instrValue.value;

214 + break;

215 + }

216 case "Debugger":

217 case "RegExpLiteral":

218 case "LoadGlobal":

@@ -510,6 +514,10 @@ export function mapInstructionValueOperands(

514 instrValue.value = fn(instrValue.value);

515 break;

516 }

517 + case "Memoize": {

518 + instrValue.value = fn(instrValue.value);

519 + break;

520 + }

521 case "Debugger":

522 case "RegExpLiteral":

523 case "LoadGlobal":

compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts

+121 -18

@@ -10,9 +10,13 @@ import {

10 Effect,

11 HIRFunction,

12 IdentifierId,

13 + Instruction,

14 Place,

15 SpreadPattern,

16 + makeInstructionId,

17 + markInstructionIds,

18 } from "../HIR";

19 +import { createTemporaryPlace } from "../HIR/HIRBuilder";

20 import { HookKind } from "../HIR/ObjectShape";

22 /*

@@ -26,8 +30,14 @@ import { HookKind } from "../HIR/ObjectShape";

30 export function dropManualMemoization(func: HIRFunction): void {

31 const hooks = new Map<IdentifierId, HookKind>();

32 const react = new Set<IdentifierId>();

33 + let hasChanges = false;

34 for (const [_, block] of func.body.blocks) {

30 - for (const instr of block.instructions) {

35 + let nextInstructions: Array<Instruction> | null = null;

36 + for (let i = 0; i < block.instructions.length; i++) {

37 + const instr = block.instructions[i]!;

38 + if (nextInstructions !== null) {

39 + nextInstructions.push(instr);

40 + }

41 switch (instr.value.kind) {

42 case "LoadGlobal": {

43 if (

@@ -71,16 +81,22 @@ export function dropManualMemoization(func: HIRFunction): void {

81 });

82 }

83 /*

74 - * TODO(gsn): Consider inlining the function passed to useMemo,

75 - * rather than just calling it directly.

76 - *

84 * Replace the hook callee with the fn arg.

85 *

86 * before:

80 - * foo = Call useMemo$2($9, $10)

87 + * $1 = LoadGlobal useMemo // load the useMemo global

88 + * $2 = FunctionExpression ... // memo function

89 + * $3 = ArrayExpression [ ... ] // deps array

90 + * $4 = Call $1 ($2, $3 ) // invoke useMemo w fn and deps

91 *

92 * after:

83 - * foo = Call $9()

93 + * $1 = LoadGlobal useMemo // load the useMemo global (dead code)

94 + * $2 = FunctionExpression ... // memo function

95 + * $3 = ArrayExpression [ ... ] // deps array (dead code)

96 + * $4 = Call $2 () // invoke the memo function itself

97 + *

98 + * Note that a later pass (InlineImmediatelyInvokedFunctionExpressions) will

99 + * inline the useMemo callback along with any other immediately invoked IIFEs.

100 */

101 if (fn.kind === "Identifier") {

102 instr.value = {

@@ -93,6 +109,46 @@ export function dropManualMemoization(func: HIRFunction): void {

109 args: [],

110 loc: instr.value.loc,

111 };

112 + if (

113 + func.env.config.enablePreserveExistingMemoizationGuarantees

114 + ) {

115 + /**

116 + * When this flag is enabled we also compile in a 'Memoize' instruction

117 + * to preserve the intended memoization boundary:

118 + *

119 + * Normal output:

120 + * $1 = LoadGlobal useMemo // load the useMemo global (dead code)

121 + * $2 = FunctionExpression ... // memo function

122 + * $3 = ArrayExpression [ ... ] // deps array (dead code)

123 + * $4 = Call $2 () // invoke the memo function itself

124 + *

125 + * Output w flag enabled:

126 + * $1 = LoadGlobal useMemo // load the useMemo global (dead code)

127 + * $2 = FunctionExpression ... // memo function

128 + * $3 = ArrayExpression [ ... ] // deps array (dead code)

129 + * $5 = Call $2 () // invoke the memo function itself

130 + * $4 = Memoize $5 // preserve memo information

131 + *

132 + * Note that we synthesize a new temporary for the call ($5) and use

133 + * the original lvalue for the result of the Memoize instruction, so that

134 + * we don't have to rewrite subsequent instructions.

135 + */

136 + const lvalue = instr.lvalue;

137 + const temp = createTemporaryPlace(func.env);

138 + instr.lvalue = { ...temp };

139 + nextInstructions =

140 + nextInstructions ?? block.instructions.slice(0, i + 1);

141 + nextInstructions.push({

142 + id: makeInstructionId(0),

143 + lvalue,

144 + value: {

145 + kind: "Memoize",

146 + value: temp,

147 + loc: instr.loc,

148 + },

149 + loc: instr.loc,

150 + });

151 + }

152 }

153 } else if (hookKind === "useCallback") {

154 const [fn] = instr.value.args as Array<

@@ -110,23 +166,63 @@ export function dropManualMemoization(func: HIRFunction): void {

166 * Instead of a Call, just alias the callback directly.

167 *

168 * before:

113 - * foo = Call useCallback$8($19)

169 + * $1 = LoadGlobal useCallback

170 + * $2 = FunctionExpression ... // the callback being memoized

171 + * $3 = ArrayExpression ... // deps array

172 + * $3 = Call $1 ( $2, $3 ) // invoke useCallback

173 *

174 * after:

116 - * foo = $19

175 + * $1 = LoadGlobal useCallback // dead code

176 + * $2 = FunctionExpression ... // the callback being memoized

177 + * $3 = ArrayExpression ... // deps array (dead code)

178 + * $3 = LoadLocal $2 // reference the function

179 */

180 if (fn.kind === "Identifier") {

119 - instr.value = {

120 - kind: "LoadLocal",

121 - place: {

122 - kind: "Identifier",

123 - identifier: fn.identifier,

124 - effect: Effect.Unknown,

125 - reactive: false,

181 + if (

182 + func.env.config.enablePreserveExistingMemoizationGuarantees

183 + ) {

184 + /**

185 + * With the flag enabled the output changes to use a Memoize instruction instead

186 + * a loadlocal to load the function expression into the original temporary:

187 + *

188 + * Normal output:

189 + * $1 = LoadGlobal useCallback // dead code

190 + * $2 = FunctionExpression ... // the callback being memoized

191 + * $3 = ArrayExpression ... // deps array (dead code)

192 + * $3 = LoadLocal $2 // reference the function

193 + *

194 + * With flag enabled:

195 + * $1 = LoadGlobal useCallback // dead code

196 + * $2 = FunctionExpression ... // the callback being memoized

197 + * $3 = ArrayExpression ... // deps array (dead code)

198 + * $3 = Memoize $2 // reference the function

199 + *

200 + * Note the s/LoadLocal/Memoize/

201 + */

202 + instr.value = {

203 + kind: "Memoize",

204 + value: {

205 + kind: "Identifier",

206 + identifier: fn.identifier,

207 + effect: Effect.Unknown,

208 + reactive: false,

209 + loc: instr.value.loc,

210 + },

211 loc: instr.value.loc,

127 - },

128 - loc: instr.value.loc,

129 - };

212 + };

213 + } else {

214 + instr.value = {

215 + kind: "LoadLocal",

216 + place: {

217 + kind: "Identifier",

218 + identifier: fn.identifier,

219 + effect: Effect.Unknown,

220 + reactive: false,

221 + loc: instr.value.loc,

222 + },

223 + loc: instr.value.loc,

224 + };

225 + }

226 }

227 }

228 }

@@ -134,5 +230,12 @@ export function dropManualMemoization(func: HIRFunction): void {

230 }

231 }

232 }

233 + if (nextInstructions !== null) {

234 + block.instructions = nextInstructions;

235 + hasChanges = true;

236 + }

237 + }

238 + if (hasChanges) {

239 + markInstructionIds(func.body);

240 }

241 }

compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts

+11

@@ -1162,6 +1162,17 @@ function inferBlock(

1162 state.alias(lvalue, instrValue.value);

1163 continue;

1164 }

1165 + case "Memoize": {

1166 + state.initialize(instrValue, {

1167 + kind: ValueKind.Frozen,

1168 + reason: new Set([ValueReason.Other]),

1169 + });

1170 + state.reference(instrValue.value, Effect.Freeze, ValueReason.Other);

1171 + const lvalue = instr.lvalue;

1172 + lvalue.effect = Effect.ConditionallyMutate;

1173 + state.alias(lvalue, instrValue.value);

1174 + continue;

1175 + }

1176 case "LoadLocal": {

1177 const lvalue = instr.lvalue;

1178 const effect =

compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts

@@ -313,6 +313,7 @@ function pruneableValue(value: InstructionValue, state: State): boolean {

313 case "StoreContext": {

314 return false;

315 }

316 + case "Memoize":

317 case "RegExpLiteral":

318 case "LoadGlobal":

319 case "ArrayExpression":

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -1628,6 +1628,10 @@ function codegenInstructionValue(

1628 );

1629 break;

1630 }

1631 + case "Memoize": {

1632 + value = codegenPlaceToExpression(cx, instrValue.value);

1633 + break;

1634 + }

1635 case "Debugger":

1636 case "DeclareLocal":

1637 case "DeclareContext":

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/InferReactiveScopeVariables.ts

+2 -1

@@ -245,7 +245,8 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {

245 case "Primitive":

246 case "NextIterableOf":

247 case "NextPropertyOf":

248 - case "Debugger": {

248 + case "Debugger":

249 + case "Memoize": {

250 return false;

251 }

252 case "UnaryExpression":

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts

@@ -471,6 +471,7 @@ function computeMemoizationInputs(

471 rvalues: [],

472 };

473 }

474 + case "Memoize":

475 case "Await":

476 case "TypeCastExpression":

477 case "NextIterableOf": {

compiler/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts

@@ -280,6 +280,11 @@ function* generateInstructionTypes(

280 break;

281 }

282

283 + case "Memoize": {

284 + yield equation(left, value.value.identifier.type);

285 + break;

286 + }

287 +

288 case "PropertyDelete":

289 case "ComputedDelete": {

290 yield equation(left, { kind: "Primitive" });

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-dont-preserve-memoization-guarantees.expect.md new

+54

@@ -0,0 +1,54 @@

1 +

2 +## Input

3 +

4 +```javascript

5 +// @enablePreserveExistingMemoizationGuarantees:false

6 +import { useMemo } from "react";

7 +import { identity, makeObject_Primitives, mutate } from "shared-runtime";

8 +

9 +function Component(props) {

10 + const object = useMemo(() => makeObject_Primitives(), []);

11 + identity(object);

12 + return object;

13 +}

14 +

15 +export const FIXTURE_ENTRYPOINT = {

16 + fn: Component,

17 + params: [{}],

18 +};

19 +

20 +```

21 +

22 +## Code

23 +

24 +```javascript

25 +// @enablePreserveExistingMemoizationGuarantees:false

26 +import { useMemo, unstable_useMemoCache as useMemoCache } from "react";

27 +import { identity, makeObject_Primitives, mutate } from "shared-runtime";

28 +

29 +function Component(props) {

30 + const $ = useMemoCache(2);

31 + let t7;

32 + let object;

33 + if ($[0] === Symbol.for("react.memo_cache_sentinel")) {

34 + t7 = makeObject_Primitives();

35 + object = t7;

36 + identity(object);

37 + $[0] = object;

38 + $[1] = t7;

39 + } else {

40 + object = $[0];

41 + t7 = $[1];

42 + }

43 + return object;

44 +}

45 +

46 +export const FIXTURE_ENTRYPOINT = {

47 + fn: Component,

48 + params: [{}],

49 +};

50 +

51 +```

52 +

53 +### Eval output

54 +(kind: ok) {"a":0,"b":"value1","c":true}

\ No newline at end of file

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-dont-preserve-memoization-guarantees.js new

+14

@@ -0,0 +1,14 @@

1 +// @enablePreserveExistingMemoizationGuarantees:false

2 +import { useMemo } from "react";

3 +import { identity, makeObject_Primitives, mutate } from "shared-runtime";

4 +

5 +function Component(props) {

6 + const object = useMemo(() => makeObject_Primitives(), []);

7 + identity(object);

8 + return object;

9 +}

10 +

11 +export const FIXTURE_ENTRYPOINT = {

12 + fn: Component,

13 + params: [{}],

14 +};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-maybe-modified-later-preserve-memoization-guarantees.expect.md new

+53

@@ -0,0 +1,53 @@

1 +

2 +## Input

3 +

4 +```javascript

5 +// @enablePreserveExistingMemoizationGuarantees

6 +import { useMemo } from "react";

7 +import { identity, makeObject_Primitives, mutate } from "shared-runtime";

8 +

9 +function Component(props) {

10 + const object = useMemo(() => makeObject_Primitives(), []);

11 + identity(object);

12 + return object;

13 +}

14 +

15 +export const FIXTURE_ENTRYPOINT = {

16 + fn: Component,

17 + params: [{}],

18 +};

19 +

20 +```

21 +

22 +## Code

23 +

24 +```javascript

25 +// @enablePreserveExistingMemoizationGuarantees

26 +import { useMemo, unstable_useMemoCache as useMemoCache } from "react";

27 +import { identity, makeObject_Primitives, mutate } from "shared-runtime";

28 +

29 +function Component(props) {

30 + const $ = useMemoCache(1);

31 + let t15;

32 + let t0;

33 + if ($[0] === Symbol.for("react.memo_cache_sentinel")) {

34 + t0 = makeObject_Primitives();

35 + $[0] = t0;

36 + } else {

37 + t0 = $[0];

38 + }

39 + t15 = t0;

40 + const object = t15;

41 + identity(object);

42 + return object;

43 +}

44 +

45 +export const FIXTURE_ENTRYPOINT = {

46 + fn: Component,

47 + params: [{}],

48 +};

49 +

50 +```

51 +

52 +### Eval output

53 +(kind: ok) {"a":0,"b":"value1","c":true}