| 1 | /** |
| 2 | * Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | * |
| 4 | * This source code is licensed under the MIT license found in the |
| 5 | * LICENSE file in the root directory of this source tree. |
| 6 | */ |
| 7 | |
| 8 | import {NodePath} from '@babel/core'; |
| 9 | import * as t from '@babel/types'; |
| 10 | import {CompilerError} from '../CompilerError'; |
| 11 | import {GeneratedSource} from '../HIR'; |
| 12 | import {ProgramContext} from './Imports'; |
| 13 | import {ExternalFunction} from '..'; |
| 14 | |
| 15 | /** |
| 16 | * Gating rewrite for function declarations which are referenced before their |
| 17 | * declaration site. |
| 18 | * |
| 19 | * ```js |
| 20 | * // original |
| 21 | * export default React.memo(Foo); |
| 22 | * function Foo() { ... } |
| 23 | * |
| 24 | * // React compiler optimized + gated |
| 25 | * import {gating} from 'myGating'; |
| 26 | * export default React.memo(Foo); |
| 27 | * const gating_result = gating(); <- inserted |
| 28 | * function Foo_optimized() {} <- inserted |
| 29 | * function Foo_unoptimized() {} <- renamed from Foo |
| 30 | * function Foo() { <- inserted function, which can be hoisted by JS engines |
| 31 | * if (gating_result) return Foo_optimized(); |
| 32 | * else return Foo_unoptimized(); |
| 33 | * } |
| 34 | * ``` |
| 35 | */ |
| 36 | function insertAdditionalFunctionDeclaration( |
| 37 | fnPath: NodePath<t.FunctionDeclaration>, |
| 38 | compiled: t.FunctionDeclaration, |
| 39 | programContext: ProgramContext, |
| 40 | gatingFunctionIdentifierName: string, |
| 41 | ): void { |
| 42 | const originalFnName = fnPath.node.id; |
| 43 | const originalFnParams = fnPath.node.params; |
| 44 | const compiledParams = fnPath.node.params; |
| 45 | /** |
| 46 | * Note that other than `export default function() {}`, all other function |
| 47 | * declarations must have a binding identifier. Since default exports cannot |
| 48 | * be referenced, it's safe to assume that all function declarations passed |
| 49 | * here will have an identifier. |
| 50 | * https://tc39.es/ecma262/multipage/ecmascript-language-functions-and-classes.html#sec-function-definitions |
| 51 | */ |
| 52 | CompilerError.invariant(originalFnName != null && compiled.id != null, { |
| 53 | reason: |
| 54 | 'Expected function declarations that are referenced elsewhere to have a named identifier', |
| 55 | loc: fnPath.node.loc ?? GeneratedSource, |
| 56 | }); |
| 57 | CompilerError.invariant(originalFnParams.length === compiledParams.length, { |
| 58 | reason: |
| 59 | 'Expected React Compiler optimized function declarations to have the same number of parameters as source', |
| 60 | loc: fnPath.node.loc ?? GeneratedSource, |
| 61 | }); |
| 62 | |
| 63 | const gatingCondition = t.identifier( |
| 64 | programContext.newUid(`${gatingFunctionIdentifierName}_result`), |
| 65 | ); |
| 66 | const unoptimizedFnName = t.identifier( |
| 67 | programContext.newUid(`${originalFnName.name}_unoptimized`), |
| 68 | ); |
| 69 | const optimizedFnName = t.identifier( |
| 70 | programContext.newUid(`${originalFnName.name}_optimized`), |
| 71 | ); |
| 72 | /** |
| 73 | * Step 1: rename existing functions |
| 74 | */ |
| 75 | compiled.id.name = optimizedFnName.name; |
| 76 | fnPath.get('id').replaceInline(unoptimizedFnName); |
| 77 | |
| 78 | /** |
| 79 | * Step 2: insert new function declaration |
| 80 | */ |
| 81 | const newParams: Array<t.Identifier | t.RestElement> = []; |
| 82 | const genNewArgs: Array<() => t.Identifier | t.SpreadElement> = []; |
| 83 | for (let i = 0; i < originalFnParams.length; i++) { |
| 84 | const argName = `arg${i}`; |
| 85 | if (originalFnParams[i].type === 'RestElement') { |
| 86 | newParams.push(t.restElement(t.identifier(argName))); |
| 87 | genNewArgs.push(() => t.spreadElement(t.identifier(argName))); |
| 88 | } else { |
| 89 | newParams.push(t.identifier(argName)); |
| 90 | genNewArgs.push(() => t.identifier(argName)); |
| 91 | } |
| 92 | } |
| 93 | // insertAfter called in reverse order of how nodes should appear in program |
| 94 | fnPath.insertAfter( |
| 95 | t.functionDeclaration( |
| 96 | originalFnName, |
| 97 | newParams, |
| 98 | t.blockStatement([ |
| 99 | t.ifStatement( |
| 100 | gatingCondition, |
| 101 | t.returnStatement( |
| 102 | t.callExpression( |
| 103 | compiled.id, |
| 104 | genNewArgs.map(fn => fn()), |
| 105 | ), |
| 106 | ), |
| 107 | t.returnStatement( |
| 108 | t.callExpression( |
| 109 | unoptimizedFnName, |
| 110 | genNewArgs.map(fn => fn()), |
| 111 | ), |
| 112 | ), |
| 113 | ), |
| 114 | ]), |
| 115 | ), |
| 116 | ); |
| 117 | fnPath.insertBefore( |
| 118 | t.variableDeclaration('const', [ |
| 119 | t.variableDeclarator( |
| 120 | gatingCondition, |
| 121 | t.callExpression(t.identifier(gatingFunctionIdentifierName), []), |
| 122 | ), |
| 123 | ]), |
| 124 | ); |
| 125 | fnPath.insertBefore(compiled); |
| 126 | } |
| 127 | export function insertGatedFunctionDeclaration( |
| 128 | fnPath: NodePath< |
| 129 | t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression |
| 130 | >, |
| 131 | compiled: |
| 132 | | t.FunctionDeclaration |
| 133 | | t.ArrowFunctionExpression |
| 134 | | t.FunctionExpression, |
| 135 | programContext: ProgramContext, |
| 136 | gating: ExternalFunction, |
| 137 | referencedBeforeDeclaration: boolean, |
| 138 | ): void { |
| 139 | const gatingImportedName = programContext.addImportSpecifier(gating).name; |
| 140 | if (referencedBeforeDeclaration && fnPath.isFunctionDeclaration()) { |
| 141 | CompilerError.invariant(compiled.type === 'FunctionDeclaration', { |
| 142 | reason: 'Expected compiled node type to match input type', |
| 143 | description: `Got ${compiled.type} but expected FunctionDeclaration`, |
| 144 | loc: fnPath.node.loc ?? GeneratedSource, |
| 145 | }); |
| 146 | insertAdditionalFunctionDeclaration( |
| 147 | fnPath, |
| 148 | compiled, |
| 149 | programContext, |
| 150 | gatingImportedName, |
| 151 | ); |
| 152 | } else { |
| 153 | const gatingExpression = t.conditionalExpression( |
| 154 | t.callExpression(t.identifier(gatingImportedName), []), |
| 155 | buildFunctionExpression(compiled), |
| 156 | buildFunctionExpression(fnPath.node), |
| 157 | ); |
| 158 | |
| 159 | /* |
| 160 | * Convert function declarations to named variables *unless* this is an |
| 161 | * `export default function ...` since `export default const ...` is |
| 162 | * not supported. For that case we fall through to replacing w the raw |
| 163 | * conditional expression |
| 164 | */ |
| 165 | if ( |
| 166 | fnPath.parentPath.node.type !== 'ExportDefaultDeclaration' && |
| 167 | fnPath.node.type === 'FunctionDeclaration' && |
| 168 | fnPath.node.id != null |
| 169 | ) { |
| 170 | fnPath.replaceWith( |
| 171 | t.variableDeclaration('const', [ |
| 172 | t.variableDeclarator(fnPath.node.id, gatingExpression), |
| 173 | ]), |
| 174 | ); |
| 175 | } else if ( |
| 176 | fnPath.parentPath.node.type === 'ExportDefaultDeclaration' && |
| 177 | fnPath.node.type !== 'ArrowFunctionExpression' && |
| 178 | fnPath.node.id != null |
| 179 | ) { |
| 180 | fnPath.insertAfter( |
| 181 | t.exportDefaultDeclaration(t.identifier(fnPath.node.id.name)), |
| 182 | ); |
| 183 | fnPath.parentPath.replaceWith( |
| 184 | t.variableDeclaration('const', [ |
| 185 | t.variableDeclarator( |
| 186 | t.identifier(fnPath.node.id.name), |
| 187 | gatingExpression, |
| 188 | ), |
| 189 | ]), |
| 190 | ); |
| 191 | } else { |
| 192 | fnPath.replaceWith(gatingExpression); |
| 193 | } |
| 194 | } |
| 195 | } |
| 196 | |
| 197 | function buildFunctionExpression( |
| 198 | node: |
| 199 | | t.FunctionDeclaration |
| 200 | | t.ArrowFunctionExpression |
| 201 | | t.FunctionExpression, |
| 202 | ): t.ArrowFunctionExpression | t.FunctionExpression { |
| 203 | if ( |
| 204 | node.type === 'ArrowFunctionExpression' || |
| 205 | node.type === 'FunctionExpression' |
| 206 | ) { |
| 207 | return node; |
| 208 | } else { |
| 209 | const fn: t.FunctionExpression = { |
| 210 | type: 'FunctionExpression', |
| 211 | async: node.async, |
| 212 | generator: node.generator, |
| 213 | loc: node.loc ?? null, |
| 214 | id: node.id ?? null, |
| 215 | params: node.params, |
| 216 | body: node.body, |
| 217 | }; |
| 218 | return fn; |
| 219 | } |
| 220 | } |