Change rules for context variable promotion
This PR updates the conditions for which variables are promoted to context variables. The previous rule was to promote any variable where the variable was reassigned in some function expression other than the function that declared the variable. Notably, this meant that we did not use context variables for variables which were captured in a function expression, but reassigned _outside_ a function expression. The new rule is more consistent: we promote any variable which is a) reassigned somewhere and b) referenced in some function expression outside of their declaring function. The implementation builds two sets of identifiers, one for each criteria, then takes the union of these two sets. ## Motivation The motivation for this change is to unblock additional validations and optimizations of function expressions. It's currently difficult to translate metadata that we infer about identifiers outside of a function expression into metadata about the identifiers within a function expression — for example to infer types within function expression bodies based on type information outside, propagate constants into functions, infer reference effects, etc. After this change, the only free variables inside function expressions will be variables that are effectively `const` - never reassigned anywhere. Thus it will be safe to renumber those identifiers to match the outer context (during EnterSSA), making it trivial to map metadata from outside the function into the function. This change also more closely models the runtime representation — any variable referenced in a function, and reassigned somewhere, would have to be compiled (ie in a JS engine) to use a context variable.
Joe Savona committed
Jun 5, 2023 at 22:11 UTC
bf631a0bb6a04244eebb0cf0fdb4fa20f2e03842
5 files changed
+115
-25
compiler/forget/packages/babel-plugin-react-forget/src/HIR/FindContextIdentifiers.ts
+38
-22
@@ -1,25 +1,26 @@
1
import type { NodePath } from "@babel/traverse";
2
import type * as t from "@babel/types";
3
import { CompilerError } from "../CompilerError";
4
+import { Set_union } from "../Utils/utils";
5
import { GeneratedSource } from "./HIR";
6
7
type FindContextIdentifierState = {
7
- inLambda: number;
8
currentLambda: Array<
9
| NodePath<t.FunctionDeclaration>
10
| NodePath<t.FunctionExpression>
11
| NodePath<t.ArrowFunctionExpression>
12
>;
13
- contextIdentifiers: Set<t.Identifier>;
13
+ reassigned: Set<t.Identifier>;
14
+ referenced: Set<t.Identifier>;
15
};
16
17
export function findContextIdentifiers(
18
func: NodePath<t.Function>
19
): Set<t.Identifier> {
20
const state: FindContextIdentifierState = {
20
- inLambda: 0,
21
currentLambda: [],
22
- contextIdentifiers: new Set(),
22
+ reassigned: new Set(),
23
+ referenced: new Set(),
24
};
25
26
func.traverse<FindContextIdentifierState>(
@@ -70,25 +71,43 @@ export function findContextIdentifiers(
71
AssignmentExpression(
72
path: NodePath<t.AssignmentExpression>,
73
state: FindContextIdentifierState
74
+ ): void {
75
+ const left = path.get("left");
76
+ handleAssignment(state.reassigned, left);
77
+ },
78
+ Identifier(
79
+ path: NodePath<t.Identifier>,
80
+ state: FindContextIdentifierState
81
): void {
82
const currentLambda = state.currentLambda.at(-1);
75
- if (currentLambda) {
76
- const left = path.get("left");
77
- handleAssignment(currentLambda, state.contextIdentifiers, left);
78
- }
83
+ if (currentLambda !== undefined)
84
+ handleIdentifier(currentLambda, state.referenced, path);
85
},
86
},
87
state
88
);
83
- return state.contextIdentifiers;
89
+ return Set_union(state.reassigned, state.referenced);
90
}
91
86
-function handleAssignment(
92
+function handleIdentifier(
93
currentLambda:
94
| NodePath<t.FunctionDeclaration>
95
| NodePath<t.FunctionExpression>
96
| NodePath<t.ArrowFunctionExpression>,
91
- contextIdentifiers: Set<t.Identifier>,
97
+ referenced: Set<t.Identifier>,
98
+ path: NodePath<t.Identifier>
99
+): void {
100
+ const name = path.node.name;
101
+ const binding = path.scope.getBinding(name);
102
+ const bindingAboveLambdaScope = currentLambda.scope.parent.getBinding(name);
103
+
104
+ if (binding != null && binding === bindingAboveLambdaScope) {
105
+ referenced.add(binding.identifier);
106
+ }
107
+}
108
+
109
+function handleAssignment(
110
+ reassigned: Set<t.Identifier>,
111
lvalPath: NodePath<t.LVal>
112
): void {
113
// Find all reassignments to identifiers declared outside of currentLambda
@@ -98,12 +117,9 @@ function handleAssignment(
117
case "Identifier": {
118
const path = lvalPath as NodePath<t.Identifier>;
119
const name = path.node.name;
101
- const ownBinding = path.scope.getBinding(name);
102
- const bindingAboveLambdaScope =
103
- currentLambda.scope.parent.getBinding(name);
104
-
105
- if (ownBinding != null && ownBinding === bindingAboveLambdaScope) {
106
- contextIdentifiers.add(ownBinding.identifier);
120
+ const binding = path.scope.getBinding(name);
121
+ if (binding != null) {
122
+ reassigned.add(binding.identifier);
123
}
124
break;
125
}
@@ -111,7 +127,7 @@ function handleAssignment(
127
const path = lvalPath as NodePath<t.ArrayPattern>;
128
for (const element of path.get("elements")) {
129
if (nonNull(element)) {
114
- handleAssignment(currentLambda, contextIdentifiers, element);
130
+ handleAssignment(reassigned, element);
131
}
132
}
133
break;
@@ -127,7 +143,7 @@ function handleAssignment(
143
valuePath.node.loc ?? GeneratedSource
144
);
145
}
130
- handleAssignment(currentLambda, contextIdentifiers, valuePath);
146
+ handleAssignment(reassigned, valuePath);
147
} else {
148
if (!property.isRestElement()) {
149
CompilerError.invariant(
@@ -135,7 +151,7 @@ function handleAssignment(
151
property.node.loc ?? GeneratedSource
152
);
153
}
138
- handleAssignment(currentLambda, contextIdentifiers, property);
154
+ handleAssignment(reassigned, property);
155
}
156
}
157
break;
@@ -143,12 +159,12 @@ function handleAssignment(
159
case "AssignmentPattern": {
160
const path = lvalPath as NodePath<t.AssignmentPattern>;
161
const left = path.get("left");
146
- handleAssignment(currentLambda, contextIdentifiers, left);
162
+ handleAssignment(reassigned, left);
163
break;
164
}
165
case "RestElement": {
166
const path = lvalPath as NodePath<t.RestElement>;
151
- handleAssignment(currentLambda, contextIdentifiers, path.get("argument"));
167
+ handleAssignment(reassigned, path.get("argument"));
168
break;
169
}
170
case "MemberExpression": {
compiler/forget/packages/babel-plugin-react-forget/src/Utils/utils.ts
+10
@@ -57,3 +57,13 @@ export function getOrInsertDefault<U, V>(
57
return defaultValue;
58
}
59
}
60
+
61
+export function Set_union<T>(a: Set<T>, b: Set<T>): Set<T> {
62
+ const union = new Set<T>();
63
+ for (const item of a) {
64
+ if (b.has(item)) {
65
+ union.add(item);
66
+ }
67
+ }
68
+ return union;
69
+}
compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/function-declaration-reassign.expect.md
+12
-3
@@ -17,15 +17,24 @@ function component() {
17
```javascript
18
import { unstable_useMemoCache as useMemoCache } from "react";
19
function component() {
20
- const $ = useMemoCache(1);
20
+ const $ = useMemoCache(2);
21
let t0;
22
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
23
- t0 = {};
23
+ t0 = function x(a) {
24
+ a.foo();
25
+ };
26
$[0] = t0;
27
} else {
28
t0 = $[0];
29
}
28
- const x = t0;
30
+ let x;
31
+ if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
32
+ x = t0;
33
+ x = {};
34
+ $[1] = x;
35
+ } else {
36
+ x = $[1];
37
+ }
38
return x;
39
}
40
compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.expect.md
new
+46
@@ -0,0 +1,46 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @debug
6
+function Component(props) {
7
+ let x = null;
8
+ const onChange = (e) => {
9
+ console.log(x);
10
+ };
11
+ x = {};
12
+ return <Foo onChange={onChange} />;
13
+}
14
+
15
+```
16
+
17
+## Code
18
+
19
+```javascript
20
+import { unstable_useMemoCache as useMemoCache } from "react"; // @debug
21
+function Component(props) {
22
+ const $ = useMemoCache(2);
23
+ let onChange;
24
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
25
+ let x;
26
+ x = null;
27
+ onChange = (e) => {
28
+ console.log(x);
29
+ };
30
+ x = {};
31
+ $[0] = onChange;
32
+ } else {
33
+ onChange = $[0];
34
+ }
35
+ let t0;
36
+ if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
37
+ t0 = <Foo onChange={onChange} />;
38
+ $[1] = t0;
39
+ } else {
40
+ t0 = $[1];
41
+ }
42
+ return t0;
43
+}
44
+
45
+```
46
+
\ No newline at end of file
compiler/forget/src/__tests__/fixtures/compiler/context-variable-reassigned-outside-of-lambda.js
new
+9
@@ -0,0 +1,9 @@
1
+// @debug
2
+function Component(props) {
3
+ let x = null;
4
+ const onChange = (e) => {
5
+ console.log(x);
6
+ };
7
+ x = {};
8
+ return <Foo onChange={onChange} />;
9
+}