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).

compiler/forget/src/Babel/BabelPlugin.ts

@@ -7,13 +7,18 @@
 
 /// <reference path="./plugin-syntax-jsx.d.ts" />
 
+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<t.FunctionDeclaration>,
+      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 },
+        });
       },
     },
   };