@samitouri / QOS-React-1 / commits / dd8b743f6a

[be] Refactor gating logic in babel plugin

`compileProgram` was getting complex, so this extracts some of the logic into smaller functions. Additionally, the `try` block now only wraps the `compileFn` generator from Pipeline, which means not accidentally catching other non-Forget errors

Lauren Tan committed Aug 11, 2023 at 16:38 UTC dd8b743f6ab9e1f3e615ed5f19d5697f59051d7e
1 file changed +160 -128
compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts
+160 -128
@@ -39,145 +39,125 @@ function hasAnyUseForgetDirectives(directives: t.Directive[]): boolean {
39 return false;
40 }
41
42 -export function compileProgram(
43 - program: NodePath<t.Program>,
42 +/**
43 + * Runs the Compiler pipeline and mutates the source AST to include the newly compiled function.
44 + * Returns a boolean denoting if the AST was mutated or not.
45 + */
46 +function compileAndInsertNewFunctionDeclaration(
47 + fnPath: NodePath<t.FunctionDeclaration>,
48 pass: CompilerPass
45 -): void {
46 - let hasForgetCompiledCode: boolean = false;
47 -
48 - function visitFn(
49 - fn: NodePath<t.FunctionDeclaration>,
50 - pass: CompilerPass
51 - ): void {
52 - try {
53 - const compiled = compileFn(fn, pass.opts.environment);
54 - if (pass.opts.noEmit === true) {
55 - return;
56 - }
57 -
58 - CompilerError.invariant(fn.node.id != null, {
59 - reason: "FunctionDeclaration must have a name",
60 - description: null,
61 - loc: fn.node.loc ?? GeneratedSource,
62 - suggestions: null,
63 - });
64 - const originalIdent = fn.node.id;
65 -
66 - if (pass.opts.gating != null) {
67 - // Rename existing function
68 - fn.node.id = addSuffix(fn.node.id, "_uncompiled");
69 -
70 - // Rename and append compiled function
71 - CompilerError.invariant(compiled.id != null, {
72 - reason: "FunctionDeclaration must produce a name",
73 - description: null,
74 - loc: fn.node.loc ?? GeneratedSource,
75 - suggestions: null,
76 - });
77 - compiled.id = addSuffix(compiled.id, "_forget");
78 - const compiledFn = fn.insertAfter(compiled)[0];
79 - compiledFn.skip();
80 -
81 - // Build and append gating test
82 - compiledFn.insertAfter(
83 - buildGatingTest({
84 - originalFnDecl: fn,
85 - compiledIdent: compiled.id,
86 - originalIdent,
87 - gating: pass.opts.gating,
88 - })
89 - );
90 - if (pass.opts.instrumentForget != null) {
91 - const instrumentFnName =
92 - pass.opts.instrumentForget.importSpecifierName;
93 - addInstrumentForget(fn, originalIdent.name, instrumentFnName);
94 - addInstrumentForget(compiledFn, originalIdent.name, instrumentFnName);
95 - }
96 - } else {
97 - fn.replaceWith(compiled);
98 - if (pass.opts.instrumentForget != null) {
99 - const instrumentFnName =
100 - pass.opts.instrumentForget.importSpecifierName;
101 - addInstrumentForget(fn, originalIdent.name, instrumentFnName);
102 - }
103 - }
104 -
105 - hasForgetCompiledCode = true;
106 - } catch (err) {
107 - if (pass.opts.logger && err) {
108 - pass.opts.logger.logEvent("err", err);
109 - }
110 - /** Always throw if the flag is enabled, otherwise we only throw if the error is critical
111 - * (eg an invariant is broken, meaning the compiler may be buggy). See
112 - * {@link CompilerError.isCritical} for mappings.
113 - * */
114 - if (
115 - pass.opts.panicOnBailout ||
116 - !(err instanceof CompilerError) ||
117 - (err instanceof CompilerError && err.isCritical())
118 - ) {
119 - throw err;
120 - } else {
121 - if (pass.opts.isDev) {
122 - log(err, pass.filename ?? null);
123 - }
49 +): boolean {
50 + let compiledFn: t.FunctionDeclaration | null;
51 + try {
52 + compiledFn = compileFn(fnPath, pass.opts.environment);
53 + } catch (err) {
54 + if (pass.opts.logger && err) {
55 + pass.opts.logger.logEvent("err", err);
56 + }
57 + /** Always throw if the flag is enabled, otherwise we only throw if the error is critical
58 + * (eg an invariant is broken, meaning the compiler may be buggy). See
59 + * {@link CompilerError.isCritical} for mappings.
60 + * */
61 + if (
62 + pass.opts.panicOnBailout ||
63 + !(err instanceof CompilerError) ||
64 + (err instanceof CompilerError && err.isCritical())
65 + ) {
66 + throw err;
67 + } else {
68 + if (pass.opts.isDev) {
69 + log(err, pass.filename ?? null);
70 }
125 - } finally {
126 - // We are generating a new FunctionDeclaration node, so we must skip over it or this
127 - // traversal will loop infinitely.
128 - fn.skip();
71 }
72 + return false;
73 }
74
132 - const visitor = {
133 - FunctionDeclaration(
134 - fn: NodePath<t.FunctionDeclaration>,
135 - pass: CompilerPass
136 - ): void {
137 - if (!shouldVisitNode(fn, pass)) {
138 - return;
139 - }
75 + if (pass.opts.noEmit === true) {
76 + return false;
77 + }
78
141 - visitFn(fn, pass);
142 - },
79 + // Sucessfully compiled
80 + if (compiledFn != null) {
81 + // We are generating a new FunctionDeclaration node, so we must skip over it or this
82 + // traversal will loop infinitely.
83 + fnPath.skip();
84
144 - ArrowFunctionExpression(
145 - fn: NodePath<t.ArrowFunctionExpression>,
146 - pass: CompilerPass
147 - ): void {
148 - if (!shouldVisitNode(fn, pass)) {
149 - return;
150 - }
85 + CompilerError.invariant(fnPath.node.id != null, {
86 + reason: "FunctionDeclaration must have a name",
87 + description: null,
88 + loc: fnPath.node.loc ?? GeneratedSource,
89 + suggestions: null,
90 + });
91 + const originalIdent = fnPath.node.id;
92 +
93 + let gatedFn = null;
94 + if (pass.opts.gating != null) {
95 + gatedFn = insertGatedFunctionDeclaration(
96 + fnPath,
97 + compiledFn,
98 + originalIdent,
99 + pass.opts.gating
100 + );
101 + } else {
102 + fnPath.replaceWith(compiledFn);
103 + }
104
152 - const loweredFn = buildFunctionDeclaration(fn);
153 - if (loweredFn instanceof CompilerErrorDetail) {
154 - const error = new CompilerError();
155 - error.pushErrorDetail(loweredFn);
105 + if (pass.opts.instrumentForget != null) {
106 + const instrumentFnName = pass.opts.instrumentForget.importSpecifierName;
107 + addInstrumentForget(fnPath, originalIdent.name, instrumentFnName);
108 + if (pass.opts.gating != null && gatedFn != null) {
109 + addInstrumentForget(gatedFn, originalIdent.name, instrumentFnName);
110 + }
111 + }
112 + return true;
113 + }
114
157 - const options = parsePluginOptions(pass.opts);
158 - if (options.logger != null) {
159 - options.logger.logEvent("err", error);
160 - }
115 + return false;
116 +}
117
162 - if (options.panicOnBailout || error.isCritical()) {
163 - throw error;
164 - } else {
165 - if (pass.opts.isDev) {
166 - log(error, pass.filename);
167 - }
168 - }
169 - return;
170 - }
118 +function insertGatedFunctionDeclaration(
119 + fnPath: NodePath<t.FunctionDeclaration>,
120 + compiled: t.FunctionDeclaration,
121 + originalIdent: t.Identifier,
122 + gating: NonNullable<PluginOptions["gating"]>
123 +): NodePath<t.FunctionDeclaration> {
124 + // Rename existing function
125 + fnPath.node.id = addSuffix(originalIdent, "_uncompiled");
126 +
127 + // Rename and append compiled function
128 + CompilerError.invariant(compiled.id != null, {
129 + reason: "FunctionDeclaration must produce a name",
130 + description: null,
131 + loc: fnPath.node.loc ?? GeneratedSource,
132 + suggestions: null,
133 + });
134 + compiled.id = addSuffix(compiled.id, "_forget");
135 + const compiledFn = fnPath.insertAfter(compiled)[0];
136 + compiledFn.skip();
137 +
138 + // Build and append gating test
139 + compiledFn.insertAfter(
140 + buildGatingTest({
141 + originalFnDecl: fnPath,
142 + compiledIdent: compiled.id,
143 + originalIdent,
144 + gating,
145 + })
146 + );
147
172 - visitFn(loweredFn, pass);
173 - },
174 - };
148 + return compiledFn;
149 +}
150
151 +export function compileProgram(
152 + program: NodePath<t.Program>,
153 + pass: CompilerPass
154 +): void {
155 const options = parsePluginOptions(pass.opts);
177 -
156 const violations = [];
157 const fileComments = pass.comments;
158 + let hasForgetCompiledCode: boolean = false;
159 let fileHasUseForgetDirective = false;
160 +
161 if (Array.isArray(fileComments)) {
162 for (const comment of fileComments) {
163 if (
@@ -241,11 +221,63 @@ export function compileProgram(
221 return;
222 }
223
244 - program.traverse(visitor, {
245 - ...pass,
246 - opts: { ...pass.opts, ...options },
247 - filename: pass.filename ?? null,
248 - });
224 + // Main traversal to compile with Forget
225 + program.traverse(
226 + {
227 + FunctionDeclaration(
228 + fn: NodePath<t.FunctionDeclaration>,
229 + pass: CompilerPass
230 + ): void {
231 + if (!shouldVisitNode(fn, pass)) {
232 + return;
233 + }
234 +
235 + hasForgetCompiledCode = compileAndInsertNewFunctionDeclaration(
236 + fn,
237 + pass
238 + );
239 + },
240 +
241 + ArrowFunctionExpression(
242 + fn: NodePath<t.ArrowFunctionExpression>,
243 + pass: CompilerPass
244 + ): void {
245 + if (!shouldVisitNode(fn, pass)) {
246 + return;
247 + }
248 +
249 + const loweredFn = buildFunctionDeclaration(fn);
250 + if (loweredFn instanceof CompilerErrorDetail) {
251 + const error = new CompilerError();
252 + error.pushErrorDetail(loweredFn);
253 +
254 + const options = parsePluginOptions(pass.opts);
255 + if (options.logger != null) {
256 + options.logger.logEvent("err", error);
257 + }
258 +
259 + if (options.panicOnBailout || error.isCritical()) {
260 + throw error;
261 + } else {
262 + if (pass.opts.isDev) {
263 + log(error, pass.filename);
264 + }
265 + }
266 + return;
267 + }
268 +
269 + hasForgetCompiledCode = compileAndInsertNewFunctionDeclaration(
270 + loweredFn,
271 + pass
272 + );
273 + },
274 + },
275 + {
276 + ...pass,
277 + opts: { ...pass.opts, ...options },
278 + filename: pass.filename ?? null,
279 + }
280 + );
281
282 // If there isn't already an import of * as React, insert it so useMemoCache doesn't
283 // throw