From e77d975b169a4e768e2efec2a506eb4a1bd1090d Mon Sep 17 00:00:00 2001
From: Lauren Tan <poteto@users.noreply.github.com>
Date: Thu, 27 Jul 2023 17:08:01 -0400
Subject: [PATCH] Flag ValidateNoSetStateInRender

This needs a bit more work before we can turn it on by default, see the next PR
for a failing test case.
---
 .../src/Entrypoint/Pipeline.ts                     | 14 ++++++++------
 .../src/HIR/Environment.ts                         | 11 +++++++++++
 .../forget/packages/snap/src/compiler-worker.ts    |  5 +++++
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts b/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
index a1abda93b9..f8f5acbed4 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
@@ -137,12 +137,14 @@ export function* run(
     validateNoRefAccessInRender(hir);
   }
 
-  const noSetStateInRenderResult = validateNoSetStateInRender(hir).unwrap();
-  yield log({
-    kind: "debug",
-    name: "ValidateNoSetStateInRender",
-    value: noSetStateInRenderResult.debug(),
-  });
+  if (env.validateNoSetStateInRender) {
+    const noSetStateInRenderResult = validateNoSetStateInRender(hir).unwrap();
+    yield log({
+      kind: "debug",
+      name: "ValidateNoSetStateInRender",
+      value: noSetStateInRenderResult.debug(),
+    });
+  }
 
   leaveSSA(hir);
   yield log({ kind: "hir", name: "LeaveSSA", value: hir });
diff --git a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts
index 16de839a47..6dce4994d1 100644
--- a/compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+++ b/compiler/forget/packages/babel-plugin-react-forget/src/HIR/Environment.ts
@@ -83,6 +83,14 @@ export type EnvironmentConfig = Partial<{
    */
   validateFrozenLambdas: boolean;
 
+  /**
+   * Validates that setState is not unconditionally called during render, as it can lead to
+   * infinite loops.
+   *
+   * Defaults to false
+   */
+  validateNoSetStateInRender: boolean;
+
   /**
    * Enable inlining of `useMemo()` function expressions so that they can be more optimally
    * compiled.
@@ -176,6 +184,7 @@ export class Environment {
   validateHooksUsage: boolean;
   validateRefAccessDuringRender: boolean;
   validateFrozenLambdas: boolean;
+  validateNoSetStateInRender: boolean;
   enableFunctionCallSignatureOptimizations: boolean;
   enableAssumeHooksFollowRulesOfReact: boolean;
   enableTreatHooksAsFunctions: boolean;
@@ -231,6 +240,8 @@ export class Environment {
     this.enableOptimizeFunctionExpressions =
       config?.enableOptimizeFunctionExpressions ?? true;
     this.assertValidMutableRanges = config?.assertValidMutableRanges ?? false;
+    this.validateNoSetStateInRender =
+      config?.validateNoSetStateInRender ?? false;
 
     this.#contextIdentifiers = contextIdentifiers;
   }
diff --git a/compiler/forget/packages/snap/src/compiler-worker.ts b/compiler/forget/packages/snap/src/compiler-worker.ts
index 67c37f0e78..e02aecb8a5 100644
--- a/compiler/forget/packages/snap/src/compiler-worker.ts
+++ b/compiler/forget/packages/snap/src/compiler-worker.ts
@@ -97,6 +97,7 @@ export async function compile(
     let enableTreatHooksAsFunctions = true;
     let disableAllMemoization = false;
     let validateRefAccessDuringRender = true;
+    let validateNoSetStateInRender = true;
     let enableEmitFreeze = null;
     let enableOptimizeFunctionExpressions = true;
     if (firstLine.indexOf("@forgetDirective") !== -1) {
@@ -132,6 +133,9 @@ export async function compile(
     if (firstLine.includes("@validateRefAccessDuringRender false")) {
       validateRefAccessDuringRender = false;
     }
+    if (firstLine.includes("@validateNoSetStateInRender false")) {
+      validateNoSetStateInRender = false;
+    }
     if (firstLine.includes("@enableOptimizeFunctionExpressions false")) {
       enableOptimizeFunctionExpressions = false;
     }
@@ -165,6 +169,7 @@ export async function compile(
         validateHooksUsage: true,
         validateRefAccessDuringRender,
         validateFrozenLambdas: true,
+        validateNoSetStateInRender,
         enableEmitFreeze,
         enableOptimizeFunctionExpressions,
         assertValidMutableRanges: true,