From cd108cd430fecda79b2d75abacbe9e5c7ac4eee8 Mon Sep 17 00:00:00 2001
From: Lauren Tan <poteto@users.noreply.github.com>
Date: Mon, 6 Feb 2023 16:47:35 -0500
Subject: [PATCH] Add enableOnlyOnUseForgetDirective flag

When this flag is enabled, Forget will only compile function declarations opted
in via the `'use forget'` directive. By default this is false.

Tested internally, see related diffs
---
 compiler/forget/src/Babel/BabelPlugin.ts | 16 ++++++++++-
 compiler/forget/src/CompilerFlags.ts     | 34 ++++++++++++++++++++++--
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/compiler/forget/src/Babel/BabelPlugin.ts b/compiler/forget/src/Babel/BabelPlugin.ts
index 12dfb82184..dc3516a419 100644
--- a/compiler/forget/src/Babel/BabelPlugin.ts
+++ b/compiler/forget/src/Babel/BabelPlugin.ts
@@ -9,6 +9,7 @@
 
 import type * as BabelCore from "@babel/core";
 import jsx from "@babel/plugin-syntax-jsx";
+import { parseCompilerFlags } from "../CompilerFlags";
 import { compile } from "../CompilerPipeline";
 
 /**
@@ -24,7 +25,20 @@ export default function ReactForgetBabelPlugin(
     inherits: jsx,
     visitor: {
       FunctionDeclaration: {
-        enter(fn, _pass) {
+        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;
           }
diff --git a/compiler/forget/src/CompilerFlags.ts b/compiler/forget/src/CompilerFlags.ts
index abd58ef284..5e41301b90 100644
--- a/compiler/forget/src/CompilerFlags.ts
+++ b/compiler/forget/src/CompilerFlags.ts
@@ -5,6 +5,36 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-type Flags = {};
+export type CompilerFlags = {
+  /**
+   * Enable to make Forget only compile functions containing the 'use forget' directive.
+   */
+  enableOnlyOnUseForgetDirective: boolean;
+};
 
-export const flags: Flags = {};
+export const defaultFlags: CompilerFlags = {
+  enableOnlyOnUseForgetDirective: false,
+} as const;
+
+export function parseCompilerFlags(obj: unknown): CompilerFlags {
+  if (obj == null || typeof obj !== "object") {
+    return defaultFlags;
+  }
+  const invalidFlags: Array<string> = [];
+  let parsedFlags: Partial<CompilerFlags> = Object.create(null);
+  for (const [key, value] of Object.entries(obj)) {
+    if (isCompilerFlag(key)) {
+      parsedFlags[key] = value;
+    } else {
+      invalidFlags.push(key);
+    }
+  }
+  if (invalidFlags.length > 0) {
+    throw new Error(`Unexpected React Forget compiler flags: ${invalidFlags}`);
+  }
+  return { ...defaultFlags, ...parsedFlags };
+}
+
+function isCompilerFlag(s: string): s is keyof typeof defaultFlags {
+  return Object.prototype.hasOwnProperty.call(defaultFlags, s);
+}