Feature to optimize function expressions
This PR adds a new feature which enables additional validation/optimization of function expressions, gated by the `enableOptimizeFunctionExpressions` feature flag. When disabled, we actually revert the changes earlier in this stack, and do all our lowering of function expressions in AnalyzeFunctions. When the feature is enabled, we incrementally process function expressions in the various compilation stages, eg InferTypes infers into function expressions, ConstantPropagation propagates constants into function expressions, etc. Because this stage optimizes function expressions, in this mode codegen uses the HIR as the source rather than the original babel node. The feature is disabled by default so it has no impact on generated code. For now i've enabled the feature on just one test to demonstrate constant propagation into a function expression.
Joe Savona committed
Jun 8, 2023 at 18:05 UTC
a8fff7cc5c67f4f8f5fbd9171081be6ef3949158
13 files changed
+194
-53
compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+13
-1
@@ -7,6 +7,7 @@
7
8
import * as t from "@babel/types";
9
import invariant from "invariant";
10
+import { ExternalFunction } from "../Entrypoint/Options";
11
import { log } from "../Utils/logger";
12
import {
13
DEFAULT_GLOBALS,
@@ -33,7 +34,6 @@ import {
34
ShapeRegistry,
35
addHook,
36
} from "./ObjectShape";
36
-import { ExternalFunction } from "../Entrypoint/Options";
37
38
export type Hook = {
39
effectKind: Effect;
@@ -150,6 +150,15 @@ export type EnvironmentConfig = Partial<{
150
* }
151
*/
152
enableEmitFreeze: ExternalFunction | null;
153
+
154
+ /**
155
+ * When enabled, function expression codegen uses a subset of the compiler pipeline
156
+ * to transform and optimize their contents. When disabled, function expression
157
+ * codegen uses the original, un-transformed function body.
158
+ *
159
+ * Defaults to false (use the un-transformed function body).
160
+ */
161
+ enableCodegenLoweredFunctionExpressions: boolean;
162
}>;
163
164
export class Environment {
@@ -165,6 +174,7 @@ export class Environment {
174
enableTreatHooksAsFunctions: boolean;
175
disableAllMemoization: boolean;
176
enableEmitFreeze: ExternalFunction | null;
177
+ enableCodegenLoweredFunctionExpressions: boolean;
178
179
#contextIdentifiers: Set<t.Identifier>;
180
@@ -208,6 +218,8 @@ export class Environment {
218
config?.enableTreatHooksAsFunctions ?? true;
219
this.disableAllMemoization = config?.disableAllMemoization ?? false;
220
this.enableEmitFreeze = config?.enableEmitFreeze ?? null;
221
+ this.enableCodegenLoweredFunctionExpressions =
222
+ config?.enableCodegenLoweredFunctionExpressions ?? false;
223
224
this.#contextIdentifiers = contextIdentifiers;
225
}
compiler/forget/packages/babel-plugin-react-forget/src/HIR/MergeConsecutiveBlocks.ts
+5
-3
@@ -31,9 +31,11 @@ import { mapOptionalFallthroughs } from "./visitors";
31
export function mergeConsecutiveBlocks(fn: HIRFunction): void {
32
const merged = new MergedBlocks();
33
for (const [, block] of fn.body.blocks) {
34
- for (const instr of block.instructions) {
35
- if (instr.value.kind === "FunctionExpression") {
36
- mergeConsecutiveBlocks(instr.value.loweredFunc);
34
+ if (fn.env.enableCodegenLoweredFunctionExpressions) {
35
+ for (const instr of block.instructions) {
36
+ if (instr.value.kind === "FunctionExpression") {
37
+ mergeConsecutiveBlocks(instr.value.loweredFunc);
38
+ }
39
}
40
}
41
compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts
+12
@@ -13,9 +13,13 @@ import {
13
Identifier,
14
isRefValueType,
15
isUseRefType,
16
+ mergeConsecutiveBlocks,
17
Place,
18
ReactiveScopeDependency,
19
} from "../HIR";
20
+import { constantPropagation } from "../Optimization";
21
+import { eliminateRedundantPhi, enterSSA } from "../SSA";
22
+import { inferTypes } from "../TypeInference";
23
import { logHIRFunction } from "../Utils/logger";
24
import { inferMutableRanges } from "./InferMutableRanges";
25
import inferReferenceEffects from "./InferReferenceEffects";
@@ -86,6 +90,14 @@ export default function analyseFunctions(func: HIRFunction): void {
90
}
91
92
function lower(func: HIRFunction): void {
93
+ if (!func.env.enableCodegenLoweredFunctionExpressions) {
94
+ mergeConsecutiveBlocks(func);
95
+ enterSSA(func);
96
+ eliminateRedundantPhi(func);
97
+ constantPropagation(func);
98
+ inferTypes(func);
99
+ }
100
+
101
analyseFunctions(func);
102
inferReferenceEffects(func, { isFunctionExpression: true });
103
inferMutableRanges(func);
compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts
+21
-10
@@ -7,6 +7,7 @@
7
8
import { isValidIdentifier } from "@babel/types";
9
import {
10
+ Environment,
11
GotoVariant,
12
HIRFunction,
13
IdentifierId,
@@ -47,7 +48,12 @@ import { eliminateRedundantPhi } from "../SSA";
48
* pass.
49
*/
50
export function constantPropagation(fn: HIRFunction): void {
50
- const haveTerminalsChanged = applyConstantPropagation(fn);
51
+ const constants: Constants = new Map();
52
+ constantPropagationImpl(fn, constants);
53
+}
54
+
55
+function constantPropagationImpl(fn: HIRFunction, constants: Constants): void {
56
+ const haveTerminalsChanged = applyConstantPropagation(fn, constants);
57
if (haveTerminalsChanged) {
58
// If terminals have changed then blocks may have become newly unreachable.
59
// Re-run minification of the graph (incl reordering instruction ids)
@@ -80,7 +86,10 @@ export function constantPropagation(fn: HIRFunction): void {
86
}
87
}
88
83
-function applyConstantPropagation(fn: HIRFunction): boolean {
89
+function applyConstantPropagation(
90
+ fn: HIRFunction,
91
+ constants: Constants
92
+): boolean {
93
// Track the set of identifiers which are used as dependencies for function expressions
94
// in order to avoid propagating these constants. This is necessary because the function
95
// itself will still reference the original value. If the dependency is propagated but the
@@ -99,8 +108,6 @@ function applyConstantPropagation(fn: HIRFunction): boolean {
108
}
109
110
let hasChanges = false;
102
-
103
- const constants: Constants = new Map();
111
for (const [, block] of fn.body.blocks) {
112
// Initialize phi values if all operands have the same known constant value.
113
// Note that this analysis uses a single-pass only, so it will never fill in
@@ -137,11 +144,13 @@ function applyConstantPropagation(fn: HIRFunction): boolean {
144
continue;
145
}
146
const instr = block.instructions[i]!;
140
- // Don't propagate constants used as function expression dependencies
141
- if (functionDependencies.has(instr.lvalue.identifier.id)) {
142
- continue;
147
+ if (!fn.env.enableCodegenLoweredFunctionExpressions) {
148
+ // Don't propagate constants used as function expression dependencies
149
+ if (functionDependencies.has(instr.lvalue.identifier.id)) {
150
+ continue;
151
+ }
152
}
144
- const value = evaluateInstruction(constants, instr);
153
+ const value = evaluateInstruction(fn.env, constants, instr);
154
if (value !== null) {
155
constants.set(instr.lvalue.identifier.id, value);
156
}
@@ -180,6 +189,7 @@ function applyConstantPropagation(fn: HIRFunction): boolean {
189
}
190
191
function evaluateInstruction(
192
+ env: Environment,
193
constants: Constants,
194
instr: Instruction
195
): Constant | null {
@@ -350,8 +360,9 @@ function evaluateInstruction(
360
return placeValue;
361
}
362
case "FunctionExpression": {
353
- // TODO: propagate constants in the outer scope into the function when traversing
354
- constantPropagation(value.loweredFunc);
363
+ if (env.enableCodegenLoweredFunctionExpressions) {
364
+ constantPropagationImpl(value.loweredFunc, constants);
365
+ }
366
return null;
367
}
368
default: {
compiler/forget/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts
+33
-1
@@ -7,6 +7,7 @@
7
8
import * as t from "@babel/types";
9
import invariant from "invariant";
10
+import { pruneUnusedLValues, pruneUnusedLabels, renameVariables } from ".";
11
import { CompilerError, ErrorSeverity } from "../CompilerError";
12
import { Environment } from "../HIR";
13
import {
@@ -30,8 +31,10 @@ import {
31
} from "../HIR/HIR";
32
import { printPlace } from "../HIR/PrintHIR";
33
import { eachPatternOperand } from "../HIR/visitors";
34
+import { deadCodeElimination } from "../Optimization";
35
import { Err, Ok, Result } from "../Utils/Result";
36
import { assertExhaustive } from "../Utils/utils";
37
+import { buildReactiveFunction } from "./BuildReactiveFunction";
38
39
export function codegenReactiveFunction(
40
fn: ReactiveFunction
@@ -957,7 +960,36 @@ function codegenInstructionValue(
960
break;
961
}
962
case "FunctionExpression": {
960
- value = t.cloneNode(instrValue.expr, true, false);
963
+ if (cx.env.enableCodegenLoweredFunctionExpressions) {
964
+ const loweredFunc = instrValue.loweredFunc;
965
+ deadCodeElimination(loweredFunc);
966
+ const reactiveFunction = buildReactiveFunction(loweredFunc);
967
+ pruneUnusedLabels(reactiveFunction);
968
+ pruneUnusedLValues(reactiveFunction);
969
+ renameVariables(reactiveFunction);
970
+ const fn = codegenReactiveFunction(reactiveFunction).unwrap();
971
+ if (instrValue.expr.type === "ArrowFunctionExpression") {
972
+ let body: t.BlockStatement | t.Expression = fn.body;
973
+ if (body.body.length === 1) {
974
+ const stmt = body.body[0]!;
975
+ if (stmt.type === "ReturnStatement" && stmt.argument != null) {
976
+ body = stmt.argument;
977
+ }
978
+ }
979
+ value = t.arrowFunctionExpression(fn.params, body, fn.async);
980
+ } else {
981
+ value = t.functionExpression(
982
+ fn.id ??
983
+ (instrValue.name != null ? t.identifier(instrValue.name) : null),
984
+ fn.params,
985
+ fn.body,
986
+ fn.generator,
987
+ fn.async
988
+ );
989
+ }
990
+ } else {
991
+ value = t.cloneNode(instrValue.expr, true, false);
992
+ }
993
break;
994
}
995
case "TaggedTemplateExpression": {
compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts
+5
-1
@@ -99,7 +99,11 @@ export function eliminateRedundantPhi(fn: HIRFunction): void {
99
rewritePlace(instr.lvalue, rewrites);
100
101
// visit function expressions on first iteration of each block
102
- if (!hasBackEdge && instr.value.kind === "FunctionExpression") {
102
+ if (
103
+ !hasBackEdge &&
104
+ instr.value.kind === "FunctionExpression" &&
105
+ fn.env.enableCodegenLoweredFunctionExpressions
106
+ ) {
107
eliminateRedundantPhi(instr.value.loweredFunc);
108
}
109
}
compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts
+13
-6
@@ -248,11 +248,15 @@ function enterSSAImpl(
248
249
if (blockId === rootEntry) {
250
// NOTE: func.context should be empty for the root function
251
- if (func.context.length !== 0) {
252
- CompilerError.invariant(
253
- `Expected function context to be empty for outer function declarations`,
254
- func.loc
255
- );
251
+ if (func.env.enableCodegenLoweredFunctionExpressions) {
252
+ if (func.context.length !== 0) {
253
+ CompilerError.invariant(
254
+ `Expected function context to be empty for outer function declarations`,
255
+ func.loc
256
+ );
257
+ }
258
+ } else {
259
+ func.context = func.context.map((p) => builder.defineContext(p));
260
}
261
func.params = func.params.map((p) => builder.definePlace(p));
262
}
@@ -261,7 +265,10 @@ function enterSSAImpl(
265
mapInstructionLValues(instr, (lvalue) => builder.definePlace(lvalue));
266
mapInstructionOperands(instr, (place) => builder.getPlace(place));
267
264
- if (instr.value.kind === "FunctionExpression") {
268
+ if (
269
+ instr.value.kind === "FunctionExpression" &&
270
+ func.env.enableCodegenLoweredFunctionExpressions
271
+ ) {
272
const loweredFunc = instr.value.loweredFunc;
273
const entry = loweredFunc.body.blocks.get(loweredFunc.body.entry)!;
274
invariant(
compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts
+8
-2
@@ -67,7 +67,11 @@ function apply(func: HIRFunction, unifier: Unifier): void {
67
}
68
const { lvalue, value } = instr;
69
lvalue.identifier.type = unifier.get(lvalue.identifier.type);
70
- if (value.kind === "FunctionExpression") {
70
+
71
+ if (
72
+ value.kind === "FunctionExpression" &&
73
+ func.env.enableCodegenLoweredFunctionExpressions
74
+ ) {
75
apply(value.loweredFunc, unifier);
76
}
77
}
@@ -247,7 +251,9 @@ function* generateInstructionTypes(
251
}
252
253
case "FunctionExpression": {
250
- yield* generate(value.loweredFunc);
254
+ if (env.enableCodegenLoweredFunctionExpressions) {
255
+ yield* generate(value.loweredFunc);
256
+ }
257
break;
258
}
259
compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md
+17
-18
@@ -20,42 +20,41 @@ function Component(props) {
20
import { unstable_useMemoCache as useMemoCache } from "react";
21
function Component(props) {
22
const $ = useMemoCache(7);
23
- const t0 = props.x;
24
- const c_0 = $[0] !== t0;
25
- let t1;
23
+ const c_0 = $[0] !== props.x;
24
+ let t0;
25
if (c_0) {
27
- t1 = foo(t0);
28
- $[0] = t0;
29
- $[1] = t1;
26
+ t0 = foo(props.x);
27
+ $[0] = props.x;
28
+ $[1] = t0;
29
} else {
31
- t1 = $[1];
30
+ t0 = $[1];
31
}
33
- const x = t1;
32
+ const x = t0;
33
const c_2 = $[2] !== props;
34
const c_3 = $[3] !== x;
36
- let t2;
35
+ let t1;
36
if (c_2 || c_3) {
38
- t2 = function () {
37
+ t1 = function () {
38
const arr = [...bar(props)];
39
return arr.at(x);
40
};
41
$[2] = props;
42
$[3] = x;
44
- $[4] = t2;
43
+ $[4] = t1;
44
} else {
46
- t2 = $[4];
45
+ t1 = $[4];
46
}
48
- const fn = t2;
47
+ const fn = t1;
48
const c_5 = $[5] !== fn;
50
- let t3;
49
+ let t2;
50
if (c_5) {
52
- t3 = fn();
51
+ t2 = fn();
52
$[5] = fn;
54
- $[6] = t3;
53
+ $[6] = t2;
54
} else {
56
- t3 = $[6];
55
+ t2 = $[6];
56
}
58
- const fnResult = t3;
57
+ const fnResult = t2;
58
return fnResult;
59
}
60
compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.expect.md
new
+43
@@ -0,0 +1,43 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @enableCodegenLoweredFunctionExpressions
6
+function Component(props) {
7
+ const x = 42;
8
+ const onEvent = () => {
9
+ console.log(x);
10
+ };
11
+ return <Foo onEvent={onEvent} />;
12
+}
13
+
14
+```
15
+
16
+## Code
17
+
18
+```javascript
19
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableCodegenLoweredFunctionExpressions
20
+function Component(props) {
21
+ const $ = useMemoCache(2);
22
+ let t0;
23
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
24
+ t0 = () => {
25
+ console.log(42);
26
+ };
27
+ $[0] = t0;
28
+ } else {
29
+ t0 = $[0];
30
+ }
31
+ const onEvent = t0;
32
+ let t1;
33
+ if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
34
+ t1 = <Foo onEvent={onEvent} />;
35
+ $[1] = t1;
36
+ } else {
37
+ t1 = $[1];
38
+ }
39
+ return t1;
40
+}
41
+
42
+```
43
+
\ No newline at end of file
compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.js
new
+8
@@ -0,0 +1,8 @@
1
+// @enableCodegenLoweredFunctionExpressions
2
+function Component(props) {
3
+ const x = 42;
4
+ const onEvent = () => {
5
+ console.log(x);
6
+ };
7
+ return <Foo onEvent={onEvent} />;
8
+}
compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md
+2
-2
@@ -38,8 +38,8 @@ function Component(props) {
38
if (c_0) {
39
const allUrls = [];
40
41
- const { media: t84, comments, urls } = post;
42
- media = t84;
41
+ const { media: t85, comments, urls } = post;
42
+ media = t85;
43
const c_3 = $[3] !== comments.length;
44
let t0;
45
if (c_3) {
compiler/forget/packages/snap/src/compiler-worker.ts
+14
-9
@@ -98,40 +98,44 @@ export async function compile(
98
let disableAllMemoization = false;
99
let validateRefAccessDuringRender = true;
100
let enableEmitFreeze = null;
101
+ let enableCodegenLoweredFunctionExpressions = false;
102
if (firstLine.indexOf("@forgetDirective") !== -1) {
103
enableOnlyOnUseForgetDirective = true;
104
}
104
- if (firstLine.indexOf("@gating") !== -1) {
105
+ if (firstLine.includes("@gating")) {
106
gating = {
107
source: "ReactForgetFeatureFlag",
108
importSpecifierName: "isForgetEnabled_Fixtures",
109
};
110
}
110
- if (firstLine.indexOf("@instrumentForget") !== -1) {
111
+ if (firstLine.includes("@instrumentForget")) {
112
instrumentForget = {
113
source: "react-forget-runtime",
114
importSpecifierName: "useRenderCounter",
115
};
116
}
116
- if (firstLine.indexOf("@panicOnBailout false") !== -1) {
117
+ if (firstLine.includes("@panicOnBailout false")) {
118
panicOnBailout = false;
119
}
119
- if (firstLine.indexOf("@memoizeJsxElements false") !== -1) {
120
+ if (firstLine.includes("@memoizeJsxElements false")) {
121
memoizeJsxElements = false;
122
}
122
- if (firstLine.indexOf("@enableAssumeHooksFollowRulesOfReact true") !== -1) {
123
+ if (firstLine.includes("@enableAssumeHooksFollowRulesOfReact true")) {
124
enableAssumeHooksFollowRulesOfReact = true;
125
}
125
- if (firstLine.indexOf("@enableTreatHooksAsFunctions false") !== -1) {
126
+ if (firstLine.includes("@enableTreatHooksAsFunctions false")) {
127
enableTreatHooksAsFunctions = false;
128
}
128
- if (firstLine.indexOf("@disableAllMemoization true") !== -1) {
129
+ if (firstLine.includes("@disableAllMemoization true")) {
130
disableAllMemoization = true;
131
}
131
- if (firstLine.indexOf("@validateRefAccessDuringRender false") !== -1) {
132
+ if (firstLine.includes("@validateRefAccessDuringRender false")) {
133
validateRefAccessDuringRender = false;
134
}
134
- if (firstLine.indexOf("@enableEmitFreeze") !== -1) {
135
+ if (firstLine.includes("@enableCodegenLoweredFunctionExpressions")) {
136
+ enableCodegenLoweredFunctionExpressions = true;
137
+ }
138
+ if (firstLine.includes("@enableEmitFreeze")) {
139
enableEmitFreeze = {
140
source: "react-forget-runtime",
141
importSpecifierName: "makeReadOnly",
@@ -162,6 +166,7 @@ export async function compile(
166
validateRefAccessDuringRender,
167
validateFrozenLambdas: true,
168
enableEmitFreeze,
169
+ enableCodegenLoweredFunctionExpressions,
170
},
171
logger: null,
172
gating,