From cadcf3de2d062d78c755a1d7672377e5933656be Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Thu, 5 Oct 2023 10:18:24 -0700
Subject: [PATCH] Reusable function for parsing config strings

Adds a helper function for parsing pragma strings to the compiler itself, and
exports it. This will be used in follow-ups to make Snap, Sprout, and Playground
all use the same pragma parser. The helper also starts from the default values,
so adopting this will also make it easy for all those places to have the same
defaults automatically.
---
 .../src/HIR/Environment.ts                    | 25 ++++++++++++++++-
 .../src/HIR/index.ts                          |  7 ++++-
 .../src/__tests__/parseConfigPragma-test.ts   | 28 +++++++++++++++++++
 .../babel-plugin-react-forget/src/index.ts    | 10 ++++++-
 4 files changed, 67 insertions(+), 3 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-forget/src/__tests__/parseConfigPragma-test.ts

diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
index fb9ef190eb..2b396e425f 100644
--- a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
@@ -212,7 +212,7 @@ export type EnvironmentConfig = {
   bailoutOnHoleyArrays: boolean;
 };
 
-const DEFAULT_ENVIRONMENT_CONFIG: Readonly<EnvironmentConfig> = {
+export const DEFAULT_ENVIRONMENT_CONFIG: Readonly<EnvironmentConfig> = {
   customHooks: null,
 
   memoizeJsxElements: true,
@@ -231,6 +231,29 @@ const DEFAULT_ENVIRONMENT_CONFIG: Readonly<EnvironmentConfig> = {
   validateRefAccessDuringRender: false,
 };
 
+export function parseConfigPragma(pragma: string): EnvironmentConfig {
+  const config = { ...DEFAULT_ENVIRONMENT_CONFIG };
+  for (const key of Object.keys(DEFAULT_ENVIRONMENT_CONFIG)) {
+    if (!isEnvironmentConfigKey(key)) {
+      continue;
+    }
+    const value = config[key];
+    if (typeof value !== "boolean") {
+      // We only support setting boolean flags via pragma strings
+      continue;
+    }
+    if (pragma.includes(`@${key}:true`)) {
+      config[key] = true as any;
+    } else if (pragma.includes(`@${key}:false`)) {
+      config[key] = false as any;
+    } else if (pragma.includes(`@${key}`)) {
+      config[key] = true as any;
+    }
+  }
+
+  return config;
+}
+
 function isEnvironmentConfigKey(key: string): key is keyof EnvironmentConfig {
   return Object.prototype.hasOwnProperty.call(DEFAULT_ENVIRONMENT_CONFIG, key);
 }
diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/index.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/index.ts
index b6f3ad631f..8de88e686c 100644
--- a/compiler/packages/babel-plugin-react-forget/src/HIR/index.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/HIR/index.ts
@@ -10,7 +10,12 @@ export { assertTerminalSuccessorsExist } from "./AssertTerminalSuccessorsExist";
 export { assertValidMutableRanges } from "./AssertValidMutableRanges";
 export { lower } from "./BuildHIR";
 export { computeDominatorTree, computePostDominatorTree } from "./Dominator";
-export { Environment, Hook } from "./Environment";
+export {
+  DEFAULT_ENVIRONMENT_CONFIG,
+  Environment,
+  Hook,
+  parseConfigPragma,
+} from "./Environment";
 export * from "./HIR";
 export {
   markInstructionIds,
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/parseConfigPragma-test.ts b/compiler/packages/babel-plugin-react-forget/src/__tests__/parseConfigPragma-test.ts
new file mode 100644
index 0000000000..efba667c40
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/parseConfigPragma-test.ts
@@ -0,0 +1,28 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { DEFAULT_ENVIRONMENT_CONFIG, parseConfigPragma } from "..";
+
+describe("parseConfigPragma()", () => {
+  it("parses flags in various forms", () => {
+    const config = parseConfigPragma(
+      "@enableForest @validateFrozenLambdas:true @memoizeJsxElements:false"
+    );
+    // Validate defaults first to make sure that the parser is getting the value from the pragma,
+    // and not just missing it and getting the default value
+    expect(DEFAULT_ENVIRONMENT_CONFIG.enableForest).toBe(false);
+    expect(DEFAULT_ENVIRONMENT_CONFIG.validateFrozenLambdas).toBe(false);
+    expect(DEFAULT_ENVIRONMENT_CONFIG.memoizeJsxElements).toBe(true);
+
+    expect(config).toEqual({
+      ...DEFAULT_ENVIRONMENT_CONFIG,
+      enableForest: true,
+      validateFrozenLambdas: true,
+      memoizeJsxElements: false,
+    });
+  });
+});
diff --git a/compiler/packages/babel-plugin-react-forget/src/index.ts b/compiler/packages/babel-plugin-react-forget/src/index.ts
index f2d79b9b4c..489ea48fb0 100644
--- a/compiler/packages/babel-plugin-react-forget/src/index.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/index.ts
@@ -19,7 +19,15 @@ export {
   parsePluginOptions,
   run,
 } from "./Entrypoint";
-export { Effect, Hook, SourceLocation, ValueKind, printHIR } from "./HIR";
+export {
+  DEFAULT_ENVIRONMENT_CONFIG,
+  Effect,
+  Hook,
+  SourceLocation,
+  ValueKind,
+  parseConfigPragma,
+  printHIR,
+} from "./HIR";
 export { printReactiveFunction } from "./ReactiveScopes";
 
 declare global {