Improve MergeConsecutiveScopes
Rewrites the core logic of MergeConsecutiveScopes to be easier to follow and fix bugs. We now do a two-pass approach: * First we iterate block instructions to identify scopes which can be merged, without actually merging the instructions themselves. * Then we iterate again, copying instructions from the block either into the new output block, or into their merged scope, as appropriate. I think the simplicity here is worth the performance cost, and we can always revisit later as necessary.
Joe Savona committed
Oct 12, 2023 at 09:47 UTC
87b1549d5c11da3370caf947231c09c39d7053da
13 files changed
+165
-124
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
+3
-3
@@ -44,8 +44,8 @@ import {
44
flattenScopesWithObjectMethods,
45
inferReactiveScopeVariables,
46
memoizeFbtOperandsInSameScope,
47
- mergeConsecutiveScopes,
47
mergeOverlappingReactiveScopes,
48
+ mergeReactiveScopesThatInvalidateTogether,
49
promoteUsedTemporaries,
50
propagateScopeDependencies,
51
pruneAllReactiveScopes,
@@ -297,10 +297,10 @@ function* runWithEnvironment(
297
});
298
299
if (env.config.enableMergeConsecutiveScopes) {
300
- mergeConsecutiveScopes(reactiveFunction);
300
+ mergeReactiveScopesThatInvalidateTogether(reactiveFunction);
301
yield log({
302
kind: "reactive",
303
- name: "MergeConsecutiveScopes",
303
+ name: "MergeReactiveScopesThatInvalidateTogether",
304
value: reactiveFunction,
305
});
306
}
compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts
+3
-1
@@ -1002,11 +1002,13 @@ export function isMutableEffect(
1002
export type ReactiveScope = {
1003
id: ScopeId;
1004
range: MutableRange;
1005
- dependencies: Set<ReactiveScopeDependency>;
1005
+ dependencies: ReactiveScopeDependencies;
1006
declarations: Map<IdentifierId, ReactiveScopeDeclaration>;
1007
reassignments: Set<Identifier>;
1008
};
1009
1010
+export type ReactiveScopeDependencies = Set<ReactiveScopeDependency>;
1011
+
1012
export type ReactiveScopeDeclaration = {
1013
identifier: Identifier;
1014
scope: ReactiveScope; // the scope in which the variable was originally declared
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts
renamed
+46
-12
@@ -15,7 +15,9 @@ import {
15
ReactiveInstruction,
16
ReactiveScope,
17
ReactiveScopeBlock,
18
+ ReactiveScopeDependencies,
19
ReactiveScopeDependency,
20
+ ReactiveStatement,
21
makeInstructionId,
22
} from "../HIR";
23
import { assertExhaustive } from "../Utils/utils";
@@ -23,6 +25,7 @@ import { printReactiveScopeSummary } from "./PrintReactiveFunction";
25
import {
26
ReactiveFunctionTransform,
27
ReactiveFunctionVisitor,
28
+ Transformed,
29
visitReactiveFunction,
30
} from "./visitors";
31
@@ -31,10 +34,13 @@ import {
34
* - Use fewer memo slots
35
* - Reduce the number of comparisons and other memoization-related instructions
36
*
34
- * This is achieved by merging consecutive reactive scopes when the two scopes
35
- * will always invalidate together. The idea is that if two scopes would always
36
- * invalidate together, it's more efficient to group the scopes together to save
37
- * on memoization overhead.
37
+ * The algorithm merges in two main cases: consecutive scopes that invalidate together
38
+ * or nested scopes that invalidate together
39
+ *
40
+ * ## Consecutive Scopes
41
+ *
42
+ * The idea is that if two consecutive scopes would always invalidate together,
43
+ * it's more efficient to group the scopes together to save on memoization overhead.
44
*
45
* This optimization is necessarily somewhat limited. First, we only merge
46
* scopes that are in the same (reactive) block, ie we don't merge across
@@ -54,15 +60,25 @@ import {
60
* may not be beneficial if the outupts of A are not guaranteed to change if its input
61
* changes, but in practice this is generally the case.
62
*
63
+ * ## Nested Scopes
64
+ *
65
+ * In this case, if an inner scope has the same dependencies as its parent, then we can
66
+ * flatten away the inner scope since it will always invalidate at the same time.
67
+ *
68
+ * Note that PropagateScopeDependencies propagates scope dependencies upwards. This ensures
69
+ * that parent scopes have the union of their own direct dependencies as well as those of
70
+ * their (transitive) children. As a result nested scopes may have the same or fewer
71
+ * dependencies than their parents, but not more dependencies. If they have fewer dependncies,
72
+ * it means that the inner scope does not always invalidate with the parent and we should not
73
+ * flatten. If they inner scope has the exact same dependencies, however, then it's always
74
+ * better to flatten.
75
*/
58
-export function mergeConsecutiveScopes(fn: ReactiveFunction): void {
76
+export function mergeReactiveScopesThatInvalidateTogether(
77
+ fn: ReactiveFunction
78
+): void {
79
const lastUsageVisitor = new FindLastUsageVisitor();
80
visitReactiveFunction(fn, lastUsageVisitor, undefined);
61
- visitReactiveFunction(
62
- fn,
63
- new Transform(lastUsageVisitor.lastUsage),
64
- undefined
65
- );
81
+ visitReactiveFunction(fn, new Transform(lastUsageVisitor.lastUsage), null);
82
}
83
84
const DEBUG: boolean = false;
@@ -85,7 +101,7 @@ class FindLastUsageVisitor extends ReactiveFunctionVisitor<void> {
101
}
102
}
103
88
-class Transform extends ReactiveFunctionTransform<void> {
104
+class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | null> {
105
lastUsage: Map<IdentifierId, InstructionId>;
106
107
constructor(lastUsage: Map<IdentifierId, InstructionId>) {
@@ -93,7 +109,25 @@ class Transform extends ReactiveFunctionTransform<void> {
109
this.lastUsage = lastUsage;
110
}
111
96
- override visitBlock(block: ReactiveBlock, state: void): void {
112
+ override transformScope(
113
+ scope: ReactiveScopeBlock,
114
+ state: ReactiveScopeDependencies | null
115
+ ): Transformed<ReactiveStatement> {
116
+ this.visitScope(scope, scope.scope.dependencies);
117
+ if (
118
+ state !== null &&
119
+ areEqualDependencies(state, scope.scope.dependencies)
120
+ ) {
121
+ return { kind: "replace-many", value: scope.instructions };
122
+ } else {
123
+ return { kind: "keep" };
124
+ }
125
+ }
126
+
127
+ override visitBlock(
128
+ block: ReactiveBlock,
129
+ state: ReactiveScopeDependencies | null
130
+ ): void {
131
// Pass 1: visit nested blocks to potentially merge their scopes
132
this.traverseBlock(block, state);
133
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/index.ts
+1
-1
@@ -18,8 +18,8 @@ export { flattenScopesWithHooks } from "./FlattenScopesWithHooks";
18
export { flattenScopesWithObjectMethods } from "./FlattenScopesWithObjectMethods";
19
export { inferReactiveScopeVariables } from "./InferReactiveScopeVariables";
20
export { memoizeFbtOperandsInSameScope } from "./MemoizeFbtOperandsInSameScope";
21
-export { mergeConsecutiveScopes } from "./MergeConsecutiveScopes";
21
export { mergeOverlappingReactiveScopes } from "./MergeOverlappingReactiveScopes";
22
+export { mergeReactiveScopesThatInvalidateTogether } from "./MergeReactiveScopesThatInvalidateTogether";
23
export { printReactiveFunction } from "./PrintReactiveFunction";
24
export { promoteUsedTemporaries } from "./PromoteUsedTemporaries";
25
export { propagateScopeDependencies } from "./PropagateScopeDependencies";
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-in-statement-break.expect.md
+2
-11
@@ -27,20 +27,11 @@ export const FIXTURE_ENTRYPOINT = {
27
```javascript
28
import { unstable_useMemoCache as useMemoCache } from "react";
29
function Component(props) {
30
- const $ = useMemoCache(4);
30
+ const $ = useMemoCache(2);
31
const c_0 = $[0] !== props.value;
32
let x;
33
if (c_0) {
34
- const c_2 = $[2] !== props.value;
35
- let t0;
36
- if (c_2) {
37
- t0 = { ...props.value };
38
- $[2] = props.value;
39
- $[3] = t0;
40
- } else {
41
- t0 = $[3];
42
- }
43
- const object = t0;
34
+ const object = { ...props.value };
35
for (const y in object) {
36
if (y === "break") {
37
break;
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/for-of-conditional-break.expect.md
+2
-9
@@ -26,18 +26,11 @@ export const FIXTURE_ENTRYPOINT = {
26
```javascript
27
import { unstable_useMemoCache as useMemoCache } from "react";
28
function Component() {
29
- const $ = useMemoCache(2);
29
+ const $ = useMemoCache(1);
30
let x;
31
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
32
x = [];
33
- let t0;
34
- if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
35
- t0 = [1, 2];
36
- $[1] = t0;
37
- } else {
38
- t0 = $[1];
39
- }
40
- for (const item of t0) {
33
+ for (const item of [1, 2]) {
34
if (item === 1) {
35
break;
36
}
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hoisting-simple-function-expression.expect.md
+9
-15
@@ -26,25 +26,19 @@ export const FIXTURE_ENTRYPOINT = {
26
```javascript
27
import { unstable_useMemoCache as useMemoCache } from "react";
28
function hoisting() {
29
- const $ = useMemoCache(2);
30
- let t1;
29
+ const $ = useMemoCache(1);
30
+ let t0;
31
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
32
const foo = () => bar();
33
- let t0;
34
- if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
35
- t0 = () => 1;
36
- $[1] = t0;
37
- } else {
38
- t0 = $[1];
39
- }
40
- const bar = t0;
41
-
42
- t1 = foo();
43
- $[0] = t1;
33
+
34
+ const bar = () => 1;
35
+
36
+ t0 = foo();
37
+ $[0] = t0;
38
} else {
45
- t1 = $[0];
39
+ t0 = $[0];
40
}
47
- return t1;
41
+ return t0;
42
}
43
44
export const FIXTURE_ENTRYPOINT = {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/inner-memo-value-not-promoted-to-outer-scope-static.expect.md
+11
-35
@@ -21,48 +21,24 @@ function Component(props) {
21
```javascript
22
import { unstable_useMemoCache as useMemoCache } from "react";
23
function Component(props) {
24
- const $ = useMemoCache(6);
25
- let t1;
26
- let T2;
24
+ const $ = useMemoCache(1);
25
let t0;
28
- let T3;
26
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
27
const count = new MaybeMutable();
28
32
- T3 = View;
33
- T2 = View;
34
- if ($[4] === Symbol.for("react.memo_cache_sentinel")) {
35
- t0 = <span>Text</span>;
36
- $[4] = t0;
37
- } else {
38
- t0 = $[4];
39
- }
40
- t1 = maybeMutate(count);
41
- $[0] = t1;
42
- $[1] = T2;
43
- $[2] = t0;
44
- $[3] = T3;
45
- } else {
46
- t1 = $[0];
47
- T2 = $[1];
48
- t0 = $[2];
49
- T3 = $[3];
50
- }
51
- let t4;
52
- if ($[5] === Symbol.for("react.memo_cache_sentinel")) {
53
- t4 = (
54
- <T3>
55
- <T2>
56
- {t0}
57
- <span>{t1}</span>
58
- </T2>
59
- </T3>
29
+ t0 = (
30
+ <View>
31
+ <View>
32
+ <span>Text</span>
33
+ <span>{maybeMutate(count)}</span>
34
+ </View>
35
+ </View>
36
);
61
- $[5] = t4;
37
+ $[0] = t0;
38
} else {
63
- t4 = $[5];
39
+ t0 = $[0];
40
}
65
- return t4;
41
+ return t0;
42
}
43
44
```
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/merge-nested-scopes-with-same-inputs.expect.md
new
+59
@@ -0,0 +1,59 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @enableMergeConsecutiveScopes
6
+function Component(props) {
7
+ // start of scope for y, depend on props.a
8
+ let y = {};
9
+
10
+ // nested scope for x, dependent on props.a
11
+ const x = {};
12
+ mutate(x, props.a);
13
+ // end of scope for x
14
+
15
+ y.a = props.a;
16
+ y.x = x;
17
+ // end of scope for y
18
+
19
+ return y;
20
+}
21
+
22
+export const FIXTURE_ENTRYPOINT = {
23
+ fn: Component,
24
+ params: [{ a: 42 }],
25
+};
26
+
27
+```
28
+
29
+## Code
30
+
31
+```javascript
32
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableMergeConsecutiveScopes
33
+function Component(props) {
34
+ const $ = useMemoCache(2);
35
+ const c_0 = $[0] !== props.a;
36
+ let y;
37
+ if (c_0) {
38
+ y = {};
39
+
40
+ const x = {};
41
+ mutate(x, props.a);
42
+
43
+ y.a = props.a;
44
+ y.x = x;
45
+ $[0] = props.a;
46
+ $[1] = y;
47
+ } else {
48
+ y = $[1];
49
+ }
50
+ return y;
51
+}
52
+
53
+export const FIXTURE_ENTRYPOINT = {
54
+ fn: Component,
55
+ params: [{ a: 42 }],
56
+};
57
+
58
+```
59
+
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/merge-nested-scopes-with-same-inputs.js
new
+21
@@ -0,0 +1,21 @@
1
+// @enableMergeConsecutiveScopes
2
+function Component(props) {
3
+ // start of scope for y, depend on props.a
4
+ let y = {};
5
+
6
+ // nested scope for x, dependent on props.a
7
+ const x = {};
8
+ mutate(x, props.a);
9
+ // end of scope for x
10
+
11
+ y.a = props.a;
12
+ y.x = x;
13
+ // end of scope for y
14
+
15
+ return y;
16
+}
17
+
18
+export const FIXTURE_ENTRYPOINT = {
19
+ fn: Component,
20
+ params: [{ a: 42 }],
21
+};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/reactive-scope-grouping.expect.md
+4
-17
@@ -25,26 +25,13 @@ export const FIXTURE_ENTRYPOINT = {
25
```javascript
26
import { unstable_useMemoCache as useMemoCache } from "react";
27
function foo() {
28
- const $ = useMemoCache(3);
28
+ const $ = useMemoCache(1);
29
let x;
30
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
31
x = {};
32
- let y;
33
- if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
34
- y = [];
35
- let t0;
36
- if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
37
- t0 = {};
38
- $[2] = t0;
39
- } else {
40
- t0 = $[2];
41
- }
42
- const z = t0;
43
- y.push(z);
44
- $[1] = y;
45
- } else {
46
- y = $[1];
47
- }
32
+ const y = [];
33
+ const z = {};
34
+ y.push(z);
35
x.y = y;
36
$[0] = x;
37
} else {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-mutate-outer-value.expect.md
+2
-11
@@ -28,7 +28,7 @@ import { unstable_useMemoCache as useMemoCache } from "react";
28
const { shallowCopy, throwErrorWithMessage } = require("shared-runtime");
29
30
function Component(props) {
31
- const $ = useMemoCache(5);
31
+ const $ = useMemoCache(3);
32
const c_0 = $[0] !== props.a;
33
let x;
34
if (c_0) {
@@ -43,16 +43,7 @@ function Component(props) {
43
}
44
x.push(t0);
45
} catch {
46
- const c_3 = $[3] !== props.a;
47
- let t1;
48
- if (c_3) {
49
- t1 = shallowCopy({ a: props.a });
50
- $[3] = props.a;
51
- $[4] = t1;
52
- } else {
53
- t1 = $[4];
54
- }
55
- x.push(t1);
46
+ x.push(shallowCopy({ a: props.a }));
47
}
48
$[0] = props.a;
49
$[1] = x;
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/type-test-field-store.expect.md
+2
-9
@@ -23,18 +23,11 @@ export const FIXTURE_ENTRYPOINT = {
23
```javascript
24
import { unstable_useMemoCache as useMemoCache } from "react";
25
function component() {
26
- const $ = useMemoCache(2);
26
+ const $ = useMemoCache(1);
27
let x;
28
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
29
x = {};
30
- let t0;
31
- if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
32
- t0 = {};
33
- $[1] = t0;
34
- } else {
35
- t0 = $[1];
36
- }
37
- const q = t0;
30
+ const q = {};
31
x.t = q;
32
$[0] = x;
33
} else {