From 2eded946f8e1d75b2e6d078cceb7589439cf5bf4 Mon Sep 17 00:00:00 2001
From: Lauren Tan <poteto@users.noreply.github.com>
Date: Fri, 2 Jun 2023 13:06:35 -0400
Subject: [PATCH] Make CompilerEntrypoint a little more generic

These are still quite Babel specific, but the interface is slightly more
generic: CompilerEntrypoint takes a non-Babel specific CompilerPass as an
argument instead of passing Babel's PluginPass directly.

In the future we can consider lowering the whole Program into HIR but that
involves a significant lift in our representation and a small amount of new
syntax to support (eg import statements), so this PR is the extent of this stack
for now
---
 compiler/forget/src/Babel/BabelPlugin.ts  |  7 ++++++-
 compiler/forget/src/CompilerEntrypoint.ts | 17 +++++++++--------
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/compiler/forget/src/Babel/BabelPlugin.ts b/compiler/forget/src/Babel/BabelPlugin.ts
index 005d0f85ff..e2d0da658d 100644
--- a/compiler/forget/src/Babel/BabelPlugin.ts
+++ b/compiler/forget/src/Babel/BabelPlugin.ts
@@ -10,6 +10,7 @@
 import jsx from "@babel/plugin-syntax-jsx";
 import type * as BabelCore from "@babel/core";
 import { compileProgram } from "../CompilerEntrypoint";
+import { parsePluginOptions } from "../CompilerOptions";
 
 /**
  * The React Forget Babel Plugin
@@ -27,7 +28,11 @@ export default function ReactForgetBabelPlugin(
       // 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(prog, pass): void {
-        compileProgram(prog, pass);
+        compileProgram(prog, {
+          opts: parsePluginOptions(pass.opts),
+          filename: pass.filename ?? null,
+          comments: pass.file.ast.comments ?? [],
+        });
       },
     },
   };
diff --git a/compiler/forget/src/CompilerEntrypoint.ts b/compiler/forget/src/CompilerEntrypoint.ts
index e64c25fe91..07d7441126 100644
--- a/compiler/forget/src/CompilerEntrypoint.ts
+++ b/compiler/forget/src/CompilerEntrypoint.ts
@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import { NodePath, PluginPass } from "@babel/core";
+import { NodePath } from "@babel/core";
 import * as t from "@babel/types";
 import {
   CompilerError,
@@ -21,9 +21,10 @@ import {
   parsePluginOptions,
 } from "./CompilerOptions";
 
-type BabelPluginPass = {
+export type CompilerPass = {
   opts: PluginOptions;
   filename: string | null;
+  comments: (t.CommentBlock | t.CommentLine)[];
 };
 
 function hasUseForgetDirective(directive: t.Directive): boolean {
@@ -41,13 +42,13 @@ function hasAnyUseForgetDirectives(directives: t.Directive[]): boolean {
 
 export function compileProgram(
   program: NodePath<t.Program>,
-  pass: PluginPass
+  pass: CompilerPass
 ): void {
   let hasForgetCompiledCode: boolean = false;
 
   function visitFn(
     fn: NodePath<t.FunctionDeclaration>,
-    pass: BabelPluginPass
+    pass: CompilerPass
   ): void {
     try {
       const compiled = compile(fn, pass.opts.environment);
@@ -147,7 +148,7 @@ export function compileProgram(
   const visitor = {
     FunctionDeclaration(
       fn: NodePath<t.FunctionDeclaration>,
-      pass: BabelPluginPass
+      pass: CompilerPass
     ): void {
       if (!shouldCompile(fn, pass)) {
         return;
@@ -158,7 +159,7 @@ export function compileProgram(
 
     ArrowFunctionExpression(
       fn: NodePath<t.ArrowFunctionExpression>,
-      pass: BabelPluginPass
+      pass: CompilerPass
     ): void {
       if (!shouldCompile(fn, pass)) {
         return;
@@ -190,7 +191,7 @@ export function compileProgram(
   const options = parsePluginOptions(pass.opts);
 
   const violations = [];
-  const fileComments = pass.file.ast.comments;
+  const fileComments = pass.comments;
   let fileHasUseForgetDirective = false;
   if (Array.isArray(fileComments)) {
     for (const comment of fileComments) {
@@ -340,7 +341,7 @@ export function compileProgram(
 
 function shouldCompile(
   fn: NodePath<t.FunctionDeclaration | t.ArrowFunctionExpression>,
-  pass: BabelPluginPass
+  pass: CompilerPass
 ): boolean {
   if (pass.opts.enableOnlyOnUseForgetDirective) {
     const body = fn.get("body");