main
ts 304 lines 9.3 KB
Raw
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 {
9 HIRFunction,
10 IdentifierId,
11 InstructionValue,
12 makeInstructionId,
13 MutableRange,
14 Place,
15 ReactiveScope,
16 } from '../HIR';
17 import {Macro} from '../HIR/Environment';
18 import {eachInstructionValueOperand} from '../HIR/visitors';
19
20 /**
21 * Whether a macro requires its arguments to be transitively inlined (eg fbt)
22 * or just avoid having the top-level values be converted to variables (eg fbt.param)
23 */
24 enum InlineLevel {
25 Transitive = 'Transitive',
26 Shallow = 'Shallow',
27 }
28 type MacroDefinition = {
29 level: InlineLevel;
30 properties: Map<string, MacroDefinition> | null;
31 };
32
33 const SHALLOW_MACRO: MacroDefinition = {
34 level: InlineLevel.Shallow,
35 properties: null,
36 };
37 const TRANSITIVE_MACRO: MacroDefinition = {
38 level: InlineLevel.Transitive,
39 properties: null,
40 };
41 const FBT_MACRO: MacroDefinition = {
42 level: InlineLevel.Transitive,
43 properties: new Map([['*', SHALLOW_MACRO]]),
44 };
45 FBT_MACRO.properties!.set('enum', FBT_MACRO);
46
47 /**
48 * This pass supports the `fbt` translation system (https://facebook.github.io/fbt/)
49 * as well as similar user-configurable macro-like APIs where it's important that
50 * the name of the function not be changed, and it's literal arguments not be turned
51 * into temporaries.
52 *
53 * ## FBT
54 *
55 * FBT provides the `<fbt>` JSX element and `fbt()` calls (which take params in the
56 * form of `<fbt:param>` children or `fbt.param()` arguments, respectively). These
57 * tags/functions have restrictions on what types of syntax may appear as props/children/
58 * arguments, notably that variable references may not appear directly — variables
59 * must always be wrapped in a `<fbt:param>` or `fbt.param()`.
60 *
61 * To ensure that Forget doesn't rewrite code to violate this restriction, we force
62 * operands to fbt tags/calls have the same scope as the tag/call itself.
63 *
64 * Note that this still allows the props/arguments of `<fbt:param>`/`fbt.param()`
65 * to be independently memoized.
66 *
67 * ## User-defined macro-like function
68 *
69 * Users can also specify their own functions to be treated similarly to fbt via the
70 * `customMacros` environment configuration. By default, user-supplied custom macros
71 * have their arguments transitively inlined.
72 */
73 export function memoizeFbtAndMacroOperandsInSameScope(
74 fn: HIRFunction,
75 ): Set<IdentifierId> {
76 const macroKinds = new Map<Macro, MacroDefinition>([
77 ...Array.from(FBT_TAGS.entries()),
78 ...(fn.env.config.customMacros ?? []).map(
79 name => [name, TRANSITIVE_MACRO] as [Macro, MacroDefinition],
80 ),
81 ]);
82 /**
83 * Forward data-flow analysis to identify all macro tags, including
84 * things like `fbt.foo.bar(...)`
85 */
86 const macroTags = populateMacroTags(fn, macroKinds);
87
88 /**
89 * Reverse data-flow analysis to merge arguments to macro *invocations*
90 * based on the kind of the macro
91 */
92 const macroValues = mergeMacroArguments(fn, macroTags, macroKinds);
93
94 return macroValues;
95 }
96
97 const FBT_TAGS: Map<string, MacroDefinition> = new Map([
98 ['fbt', FBT_MACRO],
99 ['fbt:param', SHALLOW_MACRO],
100 ['fbt:enum', FBT_MACRO],
101 ['fbt:plural', SHALLOW_MACRO],
102 ['fbs', FBT_MACRO],
103 ['fbs:param', SHALLOW_MACRO],
104 ['fbs:enum', FBT_MACRO],
105 ['fbs:plural', SHALLOW_MACRO],
106 ]);
107 export const SINGLE_CHILD_FBT_TAGS: Set<string> = new Set([
108 'fbt:param',
109 'fbs:param',
110 ]);
111
112 function populateMacroTags(
113 fn: HIRFunction,
114 macroKinds: Map<Macro, MacroDefinition>,
115 ): Map<IdentifierId, MacroDefinition> {
116 const macroTags = new Map<IdentifierId, MacroDefinition>();
117 for (const block of fn.body.blocks.values()) {
118 for (const instr of block.instructions) {
119 const {lvalue, value} = instr;
120 switch (value.kind) {
121 case 'Primitive': {
122 if (typeof value.value === 'string') {
123 const macroDefinition = macroKinds.get(value.value);
124 if (macroDefinition != null) {
125 /*
126 * We don't distinguish between tag names and strings, so record
127 * all `fbt` string literals in case they are used as a jsx tag.
128 */
129 macroTags.set(lvalue.identifier.id, macroDefinition);
130 }
131 }
132 break;
133 }
134 case 'LoadGlobal': {
135 let macroDefinition = macroKinds.get(value.binding.name);
136 if (macroDefinition != null) {
137 macroTags.set(lvalue.identifier.id, macroDefinition);
138 }
139 break;
140 }
141 case 'PropertyLoad': {
142 if (typeof value.property === 'string') {
143 const macroDefinition = macroTags.get(value.object.identifier.id);
144 if (macroDefinition != null) {
145 const propertyDefinition =
146 macroDefinition.properties != null
147 ? (macroDefinition.properties.get(value.property) ??
148 macroDefinition.properties.get('*'))
149 : null;
150 const propertyMacro = propertyDefinition ?? macroDefinition;
151 macroTags.set(lvalue.identifier.id, propertyMacro);
152 }
153 }
154 break;
155 }
156 }
157 }
158 }
159 return macroTags;
160 }
161
162 function mergeMacroArguments(
163 fn: HIRFunction,
164 macroTags: Map<IdentifierId, MacroDefinition>,
165 macroKinds: Map<Macro, MacroDefinition>,
166 ): Set<IdentifierId> {
167 const macroValues = new Set<IdentifierId>(macroTags.keys());
168 for (const block of Array.from(fn.body.blocks.values()).reverse()) {
169 for (let i = block.instructions.length - 1; i >= 0; i--) {
170 const instr = block.instructions[i]!;
171 const {lvalue, value} = instr;
172 switch (value.kind) {
173 case 'DeclareContext':
174 case 'DeclareLocal':
175 case 'Destructure':
176 case 'LoadContext':
177 case 'LoadLocal':
178 case 'PostfixUpdate':
179 case 'PrefixUpdate':
180 case 'StoreContext':
181 case 'StoreLocal': {
182 // Instructions that never need to be merged
183 break;
184 }
185 case 'CallExpression':
186 case 'MethodCall': {
187 const scope = lvalue.identifier.scope;
188 if (scope == null) {
189 continue;
190 }
191 const callee =
192 value.kind === 'CallExpression' ? value.callee : value.property;
193 const macroDefinition =
194 macroTags.get(callee.identifier.id) ??
195 macroTags.get(lvalue.identifier.id);
196 if (macroDefinition != null) {
197 visitOperands(
198 macroDefinition,
199 scope,
200 lvalue,
201 value,
202 macroValues,
203 macroTags,
204 );
205 }
206 break;
207 }
208 case 'JsxExpression': {
209 const scope = lvalue.identifier.scope;
210 if (scope == null) {
211 continue;
212 }
213 let macroDefinition;
214 if (value.tag.kind === 'Identifier') {
215 macroDefinition = macroTags.get(value.tag.identifier.id);
216 } else {
217 macroDefinition = macroKinds.get(value.tag.name);
218 }
219 macroDefinition ??= macroTags.get(lvalue.identifier.id);
220 if (macroDefinition != null) {
221 visitOperands(
222 macroDefinition,
223 scope,
224 lvalue,
225 value,
226 macroValues,
227 macroTags,
228 );
229 }
230 break;
231 }
232 default: {
233 const scope = lvalue.identifier.scope;
234 if (scope == null) {
235 continue;
236 }
237 const macroDefinition = macroTags.get(lvalue.identifier.id);
238 if (macroDefinition != null) {
239 visitOperands(
240 macroDefinition,
241 scope,
242 lvalue,
243 value,
244 macroValues,
245 macroTags,
246 );
247 }
248 break;
249 }
250 }
251 }
252 for (const phi of block.phis) {
253 const scope = phi.place.identifier.scope;
254 if (scope == null) {
255 continue;
256 }
257 const macroDefinition = macroTags.get(phi.place.identifier.id);
258 if (
259 macroDefinition == null ||
260 macroDefinition.level === InlineLevel.Shallow
261 ) {
262 continue;
263 }
264 macroValues.add(phi.place.identifier.id);
265 for (const operand of phi.operands.values()) {
266 operand.identifier.scope = scope;
267 expandFbtScopeRange(scope.range, operand.identifier.mutableRange);
268 macroTags.set(operand.identifier.id, macroDefinition);
269 macroValues.add(operand.identifier.id);
270 }
271 }
272 }
273 return macroValues;
274 }
275
276 function expandFbtScopeRange(
277 fbtRange: MutableRange,
278 extendWith: MutableRange,
279 ): void {
280 if (extendWith.start !== 0) {
281 fbtRange.start = makeInstructionId(
282 Math.min(fbtRange.start, extendWith.start),
283 );
284 }
285 }
286
287 function visitOperands(
288 macroDefinition: MacroDefinition,
289 scope: ReactiveScope,
290 lvalue: Place,
291 value: InstructionValue,
292 macroValues: Set<IdentifierId>,
293 macroTags: Map<IdentifierId, MacroDefinition>,
294 ): void {
295 macroValues.add(lvalue.identifier.id);
296 for (const operand of eachInstructionValueOperand(value)) {
297 if (macroDefinition.level === InlineLevel.Transitive) {
298 operand.identifier.scope = scope;
299 expandFbtScopeRange(scope.range, operand.identifier.mutableRange);
300 macroTags.set(operand.identifier.id, macroDefinition);
301 }
302 macroValues.add(operand.identifier.id);
303 }
304 }