diff --git a/compiler/forget/src/Babel/BabelPlugin.ts b/compiler/forget/src/Babel/BabelPlugin.ts
index 065a146700..bf5f0a86ab 100644
--- a/compiler/forget/src/Babel/BabelPlugin.ts
+++ b/compiler/forget/src/Babel/BabelPlugin.ts
@@ -7,13 +7,18 @@
///
+import * as t from "@babel/types";
import type * as BabelCore from "@babel/core";
import generate from "@babel/generator";
import jsx from "@babel/plugin-syntax-jsx";
-import { parseCompilerFlags } from "../CompilerFlags";
+import { CompilerFlags, parseCompilerFlags } from "../CompilerFlags";
import prettier from "prettier";
import { compile } from "../CompilerPipeline";
+type BabelPluginPass = {
+ opts: CompilerFlags;
+};
+
/**
* The React Forget Babel Plugin
* @param {*} _babel
@@ -22,44 +27,57 @@ import { compile } from "../CompilerPipeline";
export default function ReactForgetBabelPlugin(
_babel: typeof BabelCore
): BabelCore.PluginObj {
+ const visitor = {
+ FunctionDeclaration(
+ fn: BabelCore.NodePath,
+ pass: BabelPluginPass
+ ) {
+ if (pass.opts.enableOnlyOnUseForgetDirective) {
+ let hasUseForgetDirective = false;
+ for (const directive of fn.node.body.directives) {
+ if (directive.value.value === "use forget") {
+ hasUseForgetDirective = true;
+ break;
+ }
+ }
+ if (!hasUseForgetDirective) {
+ return;
+ }
+ }
+ if (fn.scope.getProgramParent() !== fn.scope.parent) {
+ return;
+ }
+ const ast = compile(fn);
+
+ // We are generating a new FunctionDeclaration node, so we must skip over it or this
+ // traversal will loop infinitely.
+ try {
+ fn.replaceWith(ast);
+ fn.skip();
+ } catch (err) {
+ const result = generate(ast);
+ err.message = `${err.message}\n\n${prettier.format(result.code, {
+ semi: true,
+ parser: "babel-ts",
+ })}`;
+ throw err;
+ }
+ },
+ };
+
return {
name: "react-forget",
inherits: jsx,
visitor: {
- FunctionDeclaration: {
- enter(fn, pass) {
- const flags = parseCompilerFlags(pass.opts);
- if (flags.enableOnlyOnUseForgetDirective) {
- let hasUseForgetDirective = false;
- for (const directive of fn.node.body.directives) {
- if (directive.value.value === "use forget") {
- hasUseForgetDirective = true;
- break;
- }
- }
- if (!hasUseForgetDirective) {
- return;
- }
- }
- if (fn.scope.getProgramParent() !== fn.scope.parent) {
- return;
- }
- const ast = compile(fn);
-
- // We are generating a new FunctionDeclaration node, so we must skip over it or this
- // traversal will loop infinitely.
- try {
- fn.replaceWith(ast);
- fn.skip();
- } catch (err) {
- const result = generate(ast);
- err.message = `${err.message}\n\n${prettier.format(result.code, {
- semi: true,
- parser: "babel-ts",
- })}`;
- throw err;
- }
- },
+ // Note: Babel does some "smart" merging of visitors across plugins, so even if A is inserted
+ // prior to B, if A does not have a Program visitor and B does, B will run first. We always
+ // want Forget to run true to source as possible.
+ Program(path, pass) {
+ const flags = parseCompilerFlags(pass.opts);
+ path.traverse(visitor, {
+ ...pass,
+ opts: { ...pass.opts, ...flags },
+ });
},
},
};