Option to disable memoization
Adds a `removeAllMemoization` flag that runs the entire compiler pipeline but strips out all memoization. The intent is to be able to compare (in limited use-cases) the performance of an existing app with all memoization removed, vs the performance with manual memoization, vs the performance with Forget enabled. In terms of how this works: we already strip out useMemo/useCallback since Forget is more accurate. The new option adds an extra pass that strips out all reactive scopes. Collectively this leaves ~zero memoization within components (this does leave React.memo, but close enough).
Joe Savona committed
May 31, 2023 at 08:03 UTC
af94b075c02a82ea0e3c1937036a75b9f43727d7
7 files changed
+115
-3
compiler/forget/packages/snap/src/compiler-worker.ts
+8
-3
@@ -95,6 +95,7 @@ export async function compile(
95
let memoizeJsxElements = true;
96
let enableAssumeHooksFollowRulesOfReact = false;
97
let enableTreatHooksAsFunctions = true;
98
+ let disableAllMemoization = false;
99
if (firstLine.indexOf("@forgetDirective") !== -1) {
100
enableOnlyOnUseForgetDirective = true;
101
}
@@ -128,6 +129,9 @@ export async function compile(
129
if (firstLine.indexOf("@enableTreatHooksAsFunctions false") !== -1) {
130
enableTreatHooksAsFunctions = false;
131
}
132
+ if (firstLine.indexOf("@disableAllMemoization true") !== -1) {
133
+ disableAllMemoization = true;
134
+ }
135
136
const language = parseLanguage(firstLine);
137
@@ -143,12 +147,13 @@ export async function compile(
147
},
148
],
149
]),
146
- validateHooksUsage: true,
147
- enableFunctionCallSignatureOptimizations: true,
150
enableAssumeHooksFollowRulesOfReact,
151
+ enableFunctionCallSignatureOptimizations: true,
152
+ disableAllMemoization,
153
+ enableTreatHooksAsFunctions,
154
inlineUseMemo: true,
155
memoizeJsxElements,
151
- enableTreatHooksAsFunctions,
156
+ validateHooksUsage: true,
157
},
158
logger: null,
159
gating,
compiler/forget/src/CompilerPipeline.ts
+10
@@ -40,6 +40,7 @@ import {
40
mergeOverlappingReactiveScopes,
41
promoteUsedTemporaries,
42
propagateScopeDependencies,
43
+ pruneAllReactiveScopes,
44
pruneNonEscapingScopes,
45
pruneNonReactiveDependencies,
46
pruneUnusedLabels,
@@ -159,6 +160,15 @@ export function* run(
160
value: reactiveFunction,
161
});
162
163
+ if (env.disableAllMemoization) {
164
+ pruneAllReactiveScopes(reactiveFunction);
165
+ yield log({
166
+ kind: "reactive",
167
+ name: "PruneAllReactiveScopes",
168
+ value: reactiveFunction,
169
+ });
170
+ }
171
+
172
flattenReactiveLoops(reactiveFunction);
173
yield log({
174
kind: "reactive",
compiler/forget/src/HIR/Environment.ts
+11
@@ -102,6 +102,15 @@ export type EnvironmentConfig = Partial<{
102
* Defaults to true
103
*/
104
enableTreatHooksAsFunctions: boolean;
105
+
106
+ /**
107
+ * When enabled, removes *all* memoization from the function: this includes
108
+ * removing manually added useMemo/useCallback as well as not adding Forget's
109
+ * usual useMemoCache-based memoization.
110
+ *
111
+ * Defaults to false (ie, by default memoization is enabled)
112
+ */
113
+ disableAllMemoization: boolean;
114
}>;
115
116
export class Environment {
@@ -113,6 +122,7 @@ export class Environment {
122
enableFunctionCallSignatureOptimizations: boolean;
123
enableAssumeHooksFollowRulesOfReact: boolean;
124
enableTreatHooksAsFunctions: boolean;
125
+ disableAllMemoization: boolean;
126
#contextIdentifiers: Set<t.Identifier>;
127
128
constructor(
@@ -150,6 +160,7 @@ export class Environment {
160
config?.enableAssumeHooksFollowRulesOfReact ?? false;
161
this.enableTreatHooksAsFunctions =
162
config?.enableTreatHooksAsFunctions ?? true;
163
+ this.disableAllMemoization = config?.disableAllMemoization ?? false;
164
this.#contextIdentifiers = contextIdentifiers;
165
}
166
compiler/forget/src/ReactiveScopes/PruneAllReactiveScopes.ts
new
+36
@@ -0,0 +1,36 @@
1
+/**
2
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
3
+ *
4
+ * This source code is licensed under the MIT license found in the
5
+ * LICENSE file in the root directory of this source tree.
6
+ */
7
+
8
+import {
9
+ ReactiveFunction,
10
+ ReactiveScopeBlock,
11
+ ReactiveStatement,
12
+} from "../HIR/HIR";
13
+import {
14
+ ReactiveFunctionTransform,
15
+ Transformed,
16
+ visitReactiveFunction,
17
+} from "./visitors";
18
+
19
+/**
20
+ * Removes *all* reactive scopes. Intended for experimentation only, to allow
21
+ * accurately removing memoization using the compiler pipeline to get a baseline
22
+ * for performance of a product without memoization applied.
23
+ */
24
+export function pruneAllReactiveScopes(fn: ReactiveFunction): void {
25
+ visitReactiveFunction(fn, new Transform(), undefined);
26
+}
27
+
28
+class Transform extends ReactiveFunctionTransform<void> {
29
+ override transformScope(
30
+ scopeBlock: ReactiveScopeBlock,
31
+ state: void
32
+ ): Transformed<ReactiveStatement> {
33
+ this.visitScope(scopeBlock, state);
34
+ return { kind: "replace-many", value: scopeBlock.instructions };
35
+ }
36
+}
compiler/forget/src/ReactiveScopes/index.ts
+1
@@ -18,6 +18,7 @@ export { mergeOverlappingReactiveScopes } from "./MergeOverlappingReactiveScopes
18
export { printReactiveFunction } from "./PrintReactiveFunction";
19
export { promoteUsedTemporaries } from "./PromoteUsedTemporaries";
20
export { propagateScopeDependencies } from "./PropagateScopeDependencies";
21
+export { pruneAllReactiveScopes } from "./PruneAllReactiveScopes";
22
export { pruneNonEscapingScopes } from "./PruneNonEscapingScopes";
23
export { pruneNonReactiveDependencies } from "./PruneNonReactiveDependencies";
24
export { pruneTemporaryLValues as pruneUnusedLValues } from "./PruneTemporaryLValues";
compiler/forget/src/__tests__/fixtures/compiler/remove-memoization-kitchen-sink.expect.md
new
+37
@@ -0,0 +1,37 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @disableAllMemoization true
6
+function Component(props) {
7
+ const [x, setX] = useState(() => initializeState(props));
8
+ const onChange = useCallback((e) => {
9
+ setX(e.target.value);
10
+ });
11
+ const object = { x, onChange };
12
+ return useMemo(() => {
13
+ const { x, onChange } = object;
14
+ return <input value={x} onChange={onChange} />;
15
+ }, [x]);
16
+}
17
+
18
+```
19
+
20
+## Code
21
+
22
+```javascript
23
+// @disableAllMemoization true
24
+function Component(props) {
25
+ const [x, setX] = useState(() => initializeState(props));
26
+ const onChange = (e) => {
27
+ setX(e.target.value);
28
+ };
29
+ const object = { x, onChange };
30
+
31
+ const { x: x_0, onChange: onChange_0 } = object;
32
+ const t44 = <input value={x_0} onChange={onChange_0} />;
33
+ return t44;
34
+}
35
+
36
+```
37
+
\ No newline at end of file
compiler/forget/src/__tests__/fixtures/compiler/remove-memoization-kitchen-sink.js
new
+12
@@ -0,0 +1,12 @@
1
+// @disableAllMemoization true
2
+function Component(props) {
3
+ const [x, setX] = useState(() => initializeState(props));
4
+ const onChange = useCallback((e) => {
5
+ setX(e.target.value);
6
+ });
7
+ const object = { x, onChange };
8
+ return useMemo(() => {
9
+ const { x, onChange } = object;
10
+ return <input value={x} onChange={onChange} />;
11
+ }, [x]);
12
+}