From a8fff7cc5c67f4f8f5fbd9171081be6ef3949158 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Thu, 8 Jun 2023 18:05:41 -0400 Subject: [PATCH] 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. --- .../src/HIR/Environment.ts | 14 +++++- .../src/HIR/MergeConsecutiveBlocks.ts | 8 ++-- .../src/Inference/AnalyseFunctions.ts | 12 ++++++ .../src/Optimization/ConstantPropagation.ts | 31 ++++++++----- .../ReactiveScopes/CodegenReactiveFunction.ts | 34 ++++++++++++++- .../src/SSA/EliminateRedundantPhi.ts | 6 ++- .../src/SSA/EnterSSA.ts | 19 +++++--- .../src/TypeInference/InferTypes.ts | 10 ++++- .../compiler/array-at-closure.expect.md | 35 ++++++++------- ...gation-into-function-expressions.expect.md | 43 +++++++++++++++++++ ...t-propagation-into-function-expressions.js | 8 ++++ ...ed-scope-declarations-and-locals.expect.md | 4 +- .../packages/snap/src/compiler-worker.ts | 23 ++++++---- 13 files changed, 194 insertions(+), 53 deletions(-) create mode 100644 compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.expect.md create mode 100644 compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.js diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts index 3960292b30..e050182745 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts @@ -7,6 +7,7 @@ import * as t from "@babel/types"; import invariant from "invariant"; +import { ExternalFunction } from "../Entrypoint/Options"; import { log } from "../Utils/logger"; import { DEFAULT_GLOBALS, @@ -33,7 +34,6 @@ import { ShapeRegistry, addHook, } from "./ObjectShape"; -import { ExternalFunction } from "../Entrypoint/Options"; export type Hook = { effectKind: Effect; @@ -150,6 +150,15 @@ export type EnvironmentConfig = Partial<{ * } */ enableEmitFreeze: ExternalFunction | null; + + /** + * When enabled, function expression codegen uses a subset of the compiler pipeline + * to transform and optimize their contents. When disabled, function expression + * codegen uses the original, un-transformed function body. + * + * Defaults to false (use the un-transformed function body). + */ + enableCodegenLoweredFunctionExpressions: boolean; }>; export class Environment { @@ -165,6 +174,7 @@ export class Environment { enableTreatHooksAsFunctions: boolean; disableAllMemoization: boolean; enableEmitFreeze: ExternalFunction | null; + enableCodegenLoweredFunctionExpressions: boolean; #contextIdentifiers: Set; @@ -208,6 +218,8 @@ export class Environment { config?.enableTreatHooksAsFunctions ?? true; this.disableAllMemoization = config?.disableAllMemoization ?? false; this.enableEmitFreeze = config?.enableEmitFreeze ?? null; + this.enableCodegenLoweredFunctionExpressions = + config?.enableCodegenLoweredFunctionExpressions ?? false; this.#contextIdentifiers = contextIdentifiers; } diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/MergeConsecutiveBlocks.ts b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/MergeConsecutiveBlocks.ts index 08637466a5..75523b39f7 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/MergeConsecutiveBlocks.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/MergeConsecutiveBlocks.ts @@ -31,9 +31,11 @@ import { mapOptionalFallthroughs } from "./visitors"; export function mergeConsecutiveBlocks(fn: HIRFunction): void { const merged = new MergedBlocks(); for (const [, block] of fn.body.blocks) { - for (const instr of block.instructions) { - if (instr.value.kind === "FunctionExpression") { - mergeConsecutiveBlocks(instr.value.loweredFunc); + if (fn.env.enableCodegenLoweredFunctionExpressions) { + for (const instr of block.instructions) { + if (instr.value.kind === "FunctionExpression") { + mergeConsecutiveBlocks(instr.value.loweredFunc); + } } } diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts b/compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts index 103b01dd05..e1693a89de 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/Inference/AnalyseFunctions.ts @@ -13,9 +13,13 @@ import { Identifier, isRefValueType, isUseRefType, + mergeConsecutiveBlocks, Place, ReactiveScopeDependency, } from "../HIR"; +import { constantPropagation } from "../Optimization"; +import { eliminateRedundantPhi, enterSSA } from "../SSA"; +import { inferTypes } from "../TypeInference"; import { logHIRFunction } from "../Utils/logger"; import { inferMutableRanges } from "./InferMutableRanges"; import inferReferenceEffects from "./InferReferenceEffects"; @@ -86,6 +90,14 @@ export default function analyseFunctions(func: HIRFunction): void { } function lower(func: HIRFunction): void { + if (!func.env.enableCodegenLoweredFunctionExpressions) { + mergeConsecutiveBlocks(func); + enterSSA(func); + eliminateRedundantPhi(func); + constantPropagation(func); + inferTypes(func); + } + analyseFunctions(func); inferReferenceEffects(func, { isFunctionExpression: true }); inferMutableRanges(func); diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts b/compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts index b222b025f7..c45b06b4fc 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/Optimization/ConstantPropagation.ts @@ -7,6 +7,7 @@ import { isValidIdentifier } from "@babel/types"; import { + Environment, GotoVariant, HIRFunction, IdentifierId, @@ -47,7 +48,12 @@ import { eliminateRedundantPhi } from "../SSA"; * pass. */ export function constantPropagation(fn: HIRFunction): void { - const haveTerminalsChanged = applyConstantPropagation(fn); + const constants: Constants = new Map(); + constantPropagationImpl(fn, constants); +} + +function constantPropagationImpl(fn: HIRFunction, constants: Constants): void { + const haveTerminalsChanged = applyConstantPropagation(fn, constants); if (haveTerminalsChanged) { // If terminals have changed then blocks may have become newly unreachable. // Re-run minification of the graph (incl reordering instruction ids) @@ -80,7 +86,10 @@ export function constantPropagation(fn: HIRFunction): void { } } -function applyConstantPropagation(fn: HIRFunction): boolean { +function applyConstantPropagation( + fn: HIRFunction, + constants: Constants +): boolean { // Track the set of identifiers which are used as dependencies for function expressions // in order to avoid propagating these constants. This is necessary because the function // itself will still reference the original value. If the dependency is propagated but the @@ -99,8 +108,6 @@ function applyConstantPropagation(fn: HIRFunction): boolean { } let hasChanges = false; - - const constants: Constants = new Map(); for (const [, block] of fn.body.blocks) { // Initialize phi values if all operands have the same known constant value. // Note that this analysis uses a single-pass only, so it will never fill in @@ -137,11 +144,13 @@ function applyConstantPropagation(fn: HIRFunction): boolean { continue; } const instr = block.instructions[i]!; - // Don't propagate constants used as function expression dependencies - if (functionDependencies.has(instr.lvalue.identifier.id)) { - continue; + if (!fn.env.enableCodegenLoweredFunctionExpressions) { + // Don't propagate constants used as function expression dependencies + if (functionDependencies.has(instr.lvalue.identifier.id)) { + continue; + } } - const value = evaluateInstruction(constants, instr); + const value = evaluateInstruction(fn.env, constants, instr); if (value !== null) { constants.set(instr.lvalue.identifier.id, value); } @@ -180,6 +189,7 @@ function applyConstantPropagation(fn: HIRFunction): boolean { } function evaluateInstruction( + env: Environment, constants: Constants, instr: Instruction ): Constant | null { @@ -350,8 +360,9 @@ function evaluateInstruction( return placeValue; } case "FunctionExpression": { - // TODO: propagate constants in the outer scope into the function when traversing - constantPropagation(value.loweredFunc); + if (env.enableCodegenLoweredFunctionExpressions) { + constantPropagationImpl(value.loweredFunc, constants); + } return null; } default: { diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/forget/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts index 7578866932..2a2c59ba64 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts @@ -7,6 +7,7 @@ import * as t from "@babel/types"; import invariant from "invariant"; +import { pruneUnusedLValues, pruneUnusedLabels, renameVariables } from "."; import { CompilerError, ErrorSeverity } from "../CompilerError"; import { Environment } from "../HIR"; import { @@ -30,8 +31,10 @@ import { } from "../HIR/HIR"; import { printPlace } from "../HIR/PrintHIR"; import { eachPatternOperand } from "../HIR/visitors"; +import { deadCodeElimination } from "../Optimization"; import { Err, Ok, Result } from "../Utils/Result"; import { assertExhaustive } from "../Utils/utils"; +import { buildReactiveFunction } from "./BuildReactiveFunction"; export function codegenReactiveFunction( fn: ReactiveFunction @@ -957,7 +960,36 @@ function codegenInstructionValue( break; } case "FunctionExpression": { - value = t.cloneNode(instrValue.expr, true, false); + if (cx.env.enableCodegenLoweredFunctionExpressions) { + const loweredFunc = instrValue.loweredFunc; + deadCodeElimination(loweredFunc); + const reactiveFunction = buildReactiveFunction(loweredFunc); + pruneUnusedLabels(reactiveFunction); + pruneUnusedLValues(reactiveFunction); + renameVariables(reactiveFunction); + const fn = codegenReactiveFunction(reactiveFunction).unwrap(); + if (instrValue.expr.type === "ArrowFunctionExpression") { + let body: t.BlockStatement | t.Expression = fn.body; + if (body.body.length === 1) { + const stmt = body.body[0]!; + if (stmt.type === "ReturnStatement" && stmt.argument != null) { + body = stmt.argument; + } + } + value = t.arrowFunctionExpression(fn.params, body, fn.async); + } else { + value = t.functionExpression( + fn.id ?? + (instrValue.name != null ? t.identifier(instrValue.name) : null), + fn.params, + fn.body, + fn.generator, + fn.async + ); + } + } else { + value = t.cloneNode(instrValue.expr, true, false); + } break; } case "TaggedTemplateExpression": { diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts index e701d3aee7..a333ad7f4c 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EliminateRedundantPhi.ts @@ -99,7 +99,11 @@ export function eliminateRedundantPhi(fn: HIRFunction): void { rewritePlace(instr.lvalue, rewrites); // visit function expressions on first iteration of each block - if (!hasBackEdge && instr.value.kind === "FunctionExpression") { + if ( + !hasBackEdge && + instr.value.kind === "FunctionExpression" && + fn.env.enableCodegenLoweredFunctionExpressions + ) { eliminateRedundantPhi(instr.value.loweredFunc); } } diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts index 2ecfe2e8c3..679f854f7c 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/SSA/EnterSSA.ts @@ -248,11 +248,15 @@ function enterSSAImpl( if (blockId === rootEntry) { // NOTE: func.context should be empty for the root function - if (func.context.length !== 0) { - CompilerError.invariant( - `Expected function context to be empty for outer function declarations`, - func.loc - ); + if (func.env.enableCodegenLoweredFunctionExpressions) { + if (func.context.length !== 0) { + CompilerError.invariant( + `Expected function context to be empty for outer function declarations`, + func.loc + ); + } + } else { + func.context = func.context.map((p) => builder.defineContext(p)); } func.params = func.params.map((p) => builder.definePlace(p)); } @@ -261,7 +265,10 @@ function enterSSAImpl( mapInstructionLValues(instr, (lvalue) => builder.definePlace(lvalue)); mapInstructionOperands(instr, (place) => builder.getPlace(place)); - if (instr.value.kind === "FunctionExpression") { + if ( + instr.value.kind === "FunctionExpression" && + func.env.enableCodegenLoweredFunctionExpressions + ) { const loweredFunc = instr.value.loweredFunc; const entry = loweredFunc.body.blocks.get(loweredFunc.body.entry)!; invariant( diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts index 0f42adab0b..29cbd7ed36 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts +++ b/compiler/forget/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts @@ -67,7 +67,11 @@ function apply(func: HIRFunction, unifier: Unifier): void { } const { lvalue, value } = instr; lvalue.identifier.type = unifier.get(lvalue.identifier.type); - if (value.kind === "FunctionExpression") { + + if ( + value.kind === "FunctionExpression" && + func.env.enableCodegenLoweredFunctionExpressions + ) { apply(value.loweredFunc, unifier); } } @@ -247,7 +251,9 @@ function* generateInstructionTypes( } case "FunctionExpression": { - yield* generate(value.loweredFunc); + if (env.enableCodegenLoweredFunctionExpressions) { + yield* generate(value.loweredFunc); + } break; } diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md index eb63670c41..330045c423 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md +++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-at-closure.expect.md @@ -20,42 +20,41 @@ function Component(props) { import { unstable_useMemoCache as useMemoCache } from "react"; function Component(props) { const $ = useMemoCache(7); - const t0 = props.x; - const c_0 = $[0] !== t0; - let t1; + const c_0 = $[0] !== props.x; + let t0; if (c_0) { - t1 = foo(t0); - $[0] = t0; - $[1] = t1; + t0 = foo(props.x); + $[0] = props.x; + $[1] = t0; } else { - t1 = $[1]; + t0 = $[1]; } - const x = t1; + const x = t0; const c_2 = $[2] !== props; const c_3 = $[3] !== x; - let t2; + let t1; if (c_2 || c_3) { - t2 = function () { + t1 = function () { const arr = [...bar(props)]; return arr.at(x); }; $[2] = props; $[3] = x; - $[4] = t2; + $[4] = t1; } else { - t2 = $[4]; + t1 = $[4]; } - const fn = t2; + const fn = t1; const c_5 = $[5] !== fn; - let t3; + let t2; if (c_5) { - t3 = fn(); + t2 = fn(); $[5] = fn; - $[6] = t3; + $[6] = t2; } else { - t3 = $[6]; + t2 = $[6]; } - const fnResult = t3; + const fnResult = t2; return fnResult; } diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.expect.md new file mode 100644 index 0000000000..33e55ab569 --- /dev/null +++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.expect.md @@ -0,0 +1,43 @@ + +## Input + +```javascript +// @enableCodegenLoweredFunctionExpressions +function Component(props) { + const x = 42; + const onEvent = () => { + console.log(x); + }; + return ; +} + +``` + +## Code + +```javascript +import { unstable_useMemoCache as useMemoCache } from "react"; // @enableCodegenLoweredFunctionExpressions +function Component(props) { + const $ = useMemoCache(2); + let t0; + if ($[0] === Symbol.for("react.memo_cache_sentinel")) { + t0 = () => { + console.log(42); + }; + $[0] = t0; + } else { + t0 = $[0]; + } + const onEvent = t0; + let t1; + if ($[1] === Symbol.for("react.memo_cache_sentinel")) { + t1 = ; + $[1] = t1; + } else { + t1 = $[1]; + } + return t1; +} + +``` + \ No newline at end of file diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.js b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.js new file mode 100644 index 0000000000..3688a9c681 --- /dev/null +++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/constant-propagation-into-function-expressions.js @@ -0,0 +1,8 @@ +// @enableCodegenLoweredFunctionExpressions +function Component(props) { + const x = 42; + const onEvent = () => { + console.log(x); + }; + return ; +} diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md index 536652b111..431e23cdd7 100644 --- a/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md +++ b/compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md @@ -38,8 +38,8 @@ function Component(props) { if (c_0) { const allUrls = []; - const { media: t84, comments, urls } = post; - media = t84; + const { media: t85, comments, urls } = post; + media = t85; const c_3 = $[3] !== comments.length; let t0; if (c_3) { diff --git a/compiler/forget/packages/snap/src/compiler-worker.ts b/compiler/forget/packages/snap/src/compiler-worker.ts index d4c8712ef3..41dcd97443 100644 --- a/compiler/forget/packages/snap/src/compiler-worker.ts +++ b/compiler/forget/packages/snap/src/compiler-worker.ts @@ -98,40 +98,44 @@ export async function compile( let disableAllMemoization = false; let validateRefAccessDuringRender = true; let enableEmitFreeze = null; + let enableCodegenLoweredFunctionExpressions = false; if (firstLine.indexOf("@forgetDirective") !== -1) { enableOnlyOnUseForgetDirective = true; } - if (firstLine.indexOf("@gating") !== -1) { + if (firstLine.includes("@gating")) { gating = { source: "ReactForgetFeatureFlag", importSpecifierName: "isForgetEnabled_Fixtures", }; } - if (firstLine.indexOf("@instrumentForget") !== -1) { + if (firstLine.includes("@instrumentForget")) { instrumentForget = { source: "react-forget-runtime", importSpecifierName: "useRenderCounter", }; } - if (firstLine.indexOf("@panicOnBailout false") !== -1) { + if (firstLine.includes("@panicOnBailout false")) { panicOnBailout = false; } - if (firstLine.indexOf("@memoizeJsxElements false") !== -1) { + if (firstLine.includes("@memoizeJsxElements false")) { memoizeJsxElements = false; } - if (firstLine.indexOf("@enableAssumeHooksFollowRulesOfReact true") !== -1) { + if (firstLine.includes("@enableAssumeHooksFollowRulesOfReact true")) { enableAssumeHooksFollowRulesOfReact = true; } - if (firstLine.indexOf("@enableTreatHooksAsFunctions false") !== -1) { + if (firstLine.includes("@enableTreatHooksAsFunctions false")) { enableTreatHooksAsFunctions = false; } - if (firstLine.indexOf("@disableAllMemoization true") !== -1) { + if (firstLine.includes("@disableAllMemoization true")) { disableAllMemoization = true; } - if (firstLine.indexOf("@validateRefAccessDuringRender false") !== -1) { + if (firstLine.includes("@validateRefAccessDuringRender false")) { validateRefAccessDuringRender = false; } - if (firstLine.indexOf("@enableEmitFreeze") !== -1) { + if (firstLine.includes("@enableCodegenLoweredFunctionExpressions")) { + enableCodegenLoweredFunctionExpressions = true; + } + if (firstLine.includes("@enableEmitFreeze")) { enableEmitFreeze = { source: "react-forget-runtime", importSpecifierName: "makeReadOnly", @@ -162,6 +166,7 @@ export async function compile( validateRefAccessDuringRender, validateFrozenLambdas: true, enableEmitFreeze, + enableCodegenLoweredFunctionExpressions, }, logger: null, gating,