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 }, + }); }, }, };