@samitouri / QOS-React / commits / bf631a0bb6

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 +}