Ensure Forget runs first if specified in a Babel config
This commit adds a new Program visitor to our Babel plugin which then calls our FunctionDeclaration visitor. Babel does some "smart" merging of plugin passes so so even if plugin A is inserted prior to plugin B, if A does not have a Program visitor and B does, B will run first. Note that we also can't use Forget inside of a Babel preset as plugins run _before_ presets (https://babeljs.io/docs/en/plugins/#plugin-ordering).
Lauren Tan committed
Feb 7, 2023 at 11:11 UTC
9cac0b00685f0c140a1d36e8ed0fe89ae7e227c5
1 file changed
+53
-35
compiler/forget/src/Babel/BabelPlugin.ts
+53
-35
@@ -7,13 +7,18 @@
7
8
/// <reference path="./plugin-syntax-jsx.d.ts" />
9
10
+import * as t from "@babel/types";
11
import type * as BabelCore from "@babel/core";
12
import generate from "@babel/generator";
13
import jsx from "@babel/plugin-syntax-jsx";
13
-import { parseCompilerFlags } from "../CompilerFlags";
14
+import { CompilerFlags, parseCompilerFlags } from "../CompilerFlags";
15
import prettier from "prettier";
16
import { compile } from "../CompilerPipeline";
17
18
+type BabelPluginPass = {
19
+ opts: CompilerFlags;
20
+};
21
+
22
/**
23
* The React Forget Babel Plugin
24
* @param {*} _babel
@@ -22,44 +27,57 @@ import { compile } from "../CompilerPipeline";
27
export default function ReactForgetBabelPlugin(
28
_babel: typeof BabelCore
29
): BabelCore.PluginObj {
30
+ const visitor = {
31
+ FunctionDeclaration(
32
+ fn: BabelCore.NodePath<t.FunctionDeclaration>,
33
+ pass: BabelPluginPass
34
+ ) {
35
+ if (pass.opts.enableOnlyOnUseForgetDirective) {
36
+ let hasUseForgetDirective = false;
37
+ for (const directive of fn.node.body.directives) {
38
+ if (directive.value.value === "use forget") {
39
+ hasUseForgetDirective = true;
40
+ break;
41
+ }
42
+ }
43
+ if (!hasUseForgetDirective) {
44
+ return;
45
+ }
46
+ }
47
+ if (fn.scope.getProgramParent() !== fn.scope.parent) {
48
+ return;
49
+ }
50
+ const ast = compile(fn);
51
+
52
+ // We are generating a new FunctionDeclaration node, so we must skip over it or this
53
+ // traversal will loop infinitely.
54
+ try {
55
+ fn.replaceWith(ast);
56
+ fn.skip();
57
+ } catch (err) {
58
+ const result = generate(ast);
59
+ err.message = `${err.message}\n\n${prettier.format(result.code, {
60
+ semi: true,
61
+ parser: "babel-ts",
62
+ })}`;
63
+ throw err;
64
+ }
65
+ },
66
+ };
67
+
68
return {
69
name: "react-forget",
70
inherits: jsx,
71
visitor: {
29
- FunctionDeclaration: {
30
- enter(fn, pass) {
31
- const flags = parseCompilerFlags(pass.opts);
32
- if (flags.enableOnlyOnUseForgetDirective) {
33
- let hasUseForgetDirective = false;
34
- for (const directive of fn.node.body.directives) {
35
- if (directive.value.value === "use forget") {
36
- hasUseForgetDirective = true;
37
- break;
38
- }
39
- }
40
- if (!hasUseForgetDirective) {
41
- return;
42
- }
43
- }
44
- if (fn.scope.getProgramParent() !== fn.scope.parent) {
45
- return;
46
- }
47
- const ast = compile(fn);
48
-
49
- // We are generating a new FunctionDeclaration node, so we must skip over it or this
50
- // traversal will loop infinitely.
51
- try {
52
- fn.replaceWith(ast);
53
- fn.skip();
54
- } catch (err) {
55
- const result = generate(ast);
56
- err.message = `${err.message}\n\n${prettier.format(result.code, {
57
- semi: true,
58
- parser: "babel-ts",
59
- })}`;
60
- throw err;
61
- }
62
- },
72
+ // Note: Babel does some "smart" merging of visitors across plugins, so even if A is inserted
73
+ // prior to B, if A does not have a Program visitor and B does, B will run first. We always
74
+ // want Forget to run true to source as possible.
75
+ Program(path, pass) {
76
+ const flags = parseCompilerFlags(pass.opts);
77
+ path.traverse(visitor, {
78
+ ...pass,
79
+ opts: { ...pass.opts, ...flags },
80
+ });
81
},
82
},
83
};